Attention~AK

Multi Head Attention (MHA)

MHA MHA is the standard attention mechanism used in the original Transformer architecture. Instead of computing one attention over the entire representation space, MHA projects the input into several low-dimensional attention heads, performs attention independently in each head and then combines all output representations in to one final representation.

The important idea behind using multiple heads is that different heads can learn to attend according to different relationships (ex: syntactic dependencies, short and long-range context etc.). These is not guaranteed to happen but MHA provides separate learned subspace in which such behavior may emerge.

Each head $h$ learns three projection matrices:

W_{h}^{Q}, W_{h}^{K}, W_{h}^{V}

which produce:

Q_{h} = X W_{h}^{Q}, K_{h} = X W_{h}^{K}, V_{h} = X W_{h}^{V}

then we apply attention operation for each head $h$ as:

head_{h} = softmax (\frac{Q _{h} K _{h}^{⊤}}{d _{k}}) V_{h}

After computing all $N$ heads, we concatenate their outputs and project it using $W_{o}$ :

M H A (X) = concatenation of all attention heads [head_{1}; head_{2}; head_{3}; \dots, head_{N}] W_{o}

Implementation:

class MultiHeadAttentionBlock(nn.Module):
    
    def __init__(self, d_model: int, heads: int, dropout: float, rope = None):
        super().__init__()
        self.d_model = d_model
        self.h = heads
        
        assert self.d_model % self.h == 0, "d_model should be divisible by heads (h)"
        
        self.d_k = d_model // heads
        self.w_q = nn.Linear(d_model, d_model) # Wq
        self.w_k = nn.Linear(d_model, d_model) # Wk
        self.w_v = nn.Linear(d_model, d_model) # Wv
        
        self.w_o = nn.Linear(d_model, d_model) # Wo
        
        self.rope = rope
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, q, k, v, mask):
        query = self.w_q(q)  ## (B,S,d) --> (B,S,d)
        key = self.w_k(k)    ## (B,S,d) --> (B,S,d)
        value = self.w_v(v)  ## (B,S,d) --> (B,S,d)
        
        ## (B,S,d) --> (B,S,h,d_k) --> (B,h,S,d_k)
        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(1, 2)
        key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1, 2) 
        value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(1, 2) 
        
        
        if self.rope is not None:
            query = self.rope(query)
            key = self.rope(key)
        
        x, _ = MultiHeadAttentionBlock.attention(query, key, value, mask, self.dropout)
        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.d_model) ## (B,h,S,d_k) ---> (B,S,h,d_k) --> (B,S,d)
        
        return self.w_o(x) ## (B,S,d) --> (B,S,d)
    
    @staticmethod
    def attention(query, key, value, mask, dropout: nn.Dropout):
        d_k = query.shape[-1]
        
        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k) ## (B,h,S,d_k) * (B,h,d_k,S) --> (B,h,S,S)
        if mask is not None:
            # dtype-aware: a hardcoded -1e9 becomes -inf in fp16/AMP and yields NaNs.
            attention_scores.masked_fill_(mask == 0, torch.finfo(attention_scores.dtype).min)
        
        attention_scores = attention_scores.softmax(dim=-1) ## (B,h,S,S)
        
        if dropout is not None:
            attention_scores = dropout(attention_scores)
        
        return (attention_scores @ value), attention_scores ## (B,h,S,S) * (B,h,S,d_k) --> (B,h,S,d_k)

Gains:

Parallelization
Long distance dependency

Drawbacks:

Quadratic computation in self-attention $O (N^{2})$
Redundant heads
Slow inference - Large KV cache KV cache

Longformer Attention

To address the quadratic computation limitation of standard Transformer self-attention LongFormer architecture was proposed. LongFormer replaces full attention with sparse attention pattern. Instead of allowing every token to attend every other token, each token only attends to a carefully chosen subset of tokens.

Sliding-Window Attention

Sliding_Window

We introduce a fixed window of size $w$ . Each token only attends to tokens within a fixed local neighborhood around itself. So a token attends to $\approx w /2$ on either side. So instead of computing attention over $N$ tokens we compute it for $\approx w$ tokens.

This brings down the complexity to $O (w N)$

Dilating Sliding-Window Attention

Dilating_Sliding_Window

Instead of selecting consecutive tokens in the local window, the model introduces a dilation i.e. tokens selected by the window are separated by gaps.

Earlier, with sliding window we attended to tokens for $w = 2$ like:

i - 2, i - 1, i, i + 1, i + 2

now with dilation $d = 2$ it becomes:

i - 4, i - 2, i, i + 2, i + 4

The number of attended tokens can remain approximately the same but the window now covers a large range of the sequence hence the model can access a wider contextual region without increasing the number of tokens attended to by each query.

Global Attention

Global_Attention

Both Dilated and Sliding-Window can’t provide arbitrary global connectivity hence Longformer introduces a small amount of global attending tokens.

A globally attending token can:

attend to every token in the sequence (row)
be attended to by every token in the sequence (column)

For $g$ global tokens the computational complexity is $O (g N)$

The overall complexity of Longformer is $O (N w + N g)$

Drawbacks:

Potential for information loss
Task specific global attention studied only for QA, Conference Resolution, Document Classification
May not have significant improvement in tasks that require shorter local context.

Multi Query Attention (MQA)

Multi_Query_Attention MQA is a modification of MHA designed primarily to make Transformer inference more memory and bandwidth efficient. Instead of having a separate key and value projection for every head, all query heads share a single key and value head.

For a generated sequence of length $n$ , with $h$ heads and each head having dimension $d_{h}$ , the KV cache will contain $\approx 2 nh d_{h}$ elements per layer for MHA while for MQA it will be $\approx 2 n d_{h}$ elements per layer. So we are reducing the KV cache size by $\approx h$ when compared to MHA.

Gains:

Much lower memory-bandwidth requirements in the incremental setting

Drawbacks:

Reduced quality and capacity due to shared K & V
Address only reloading key-value during inference
Train a model just for faster inference?

Grouped Query Attention (GQA)

Grouped_Query_Attention

GQA is an intermediate form of attention between MHA and MQA. Its purpose is to reduce the memory and bandwidth cost of key-value cache during inference while attending better representational quality than MQA

The central idea is to keep multiple query heads but divide them into group and let each group of queries share only pair of key and value head. Thus, if there are $H$ query heads and $G$ KV groups, each group will contain $H / G$ query heads. The KV cache comes out to be $\approx 2 n G d_{h}$ instead of $\approx 2 n H d_{h}$ (MHA) so we are reducing the KV cache by a factor of $\approx H / G$

Implementation

class GroupedQueryAttentionBlock(nn.Module):
    
    def __init__(self, d_model: int, q_heads: int, kv_heads: int, dropout: float):
        super().__init__()
        assert d_model % q_heads == 0, "d_model should be divisible by q_heads"
        assert q_heads % kv_heads == 0, "q_heads should be divisible by kv_heads"
        
        self.d_model = d_model
        self.q_heads = q_heads
        self.kv_heads = kv_heads
        self.d_k = d_model // q_heads
        self.group_size = q_heads // kv_heads
        
        self.w_q = nn.Linear(d_model, d_model) # Wq
        self.w_k = nn.Linear(d_model, self.kv_heads*self.d_k) # Wk
        self.w_v = nn.Linear(d_model, self.kv_heads*self.d_k) # Wv
        self.w_o = nn.Linear(d_model, d_model) # Wo
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, q, k, v, mask):
        query = self.w_q(q)  ## (B,S,d) --> (B,S,d)
        key = self.w_k(k)    ## (B,S,d) --> (B,S,d_k*kv_heads)
        value = self.w_v(v)  ## (B,S,d) --> (B,S,d_k*kv_heads)
        
        ## (B,S,d) --> (B,S,h,d_k) --> (B,h,S,d_k)
        query = query.view(query.shape[0], query.shape[1], self.q_heads, self.d_k).transpose(1, 2)
        
        ## (B,S,d_k*kv_heads) --> (B,S,kv_heads,d_k) --> (B,kv_heads,S,d_k)
        key = key.view(key.shape[0], key.shape[1], self.kv_heads, self.d_k).transpose(1, 2) 
        value = value.view(value.shape[0], value.shape[1], self.kv_heads, self.d_k).transpose(1, 2)
        
        key = key.repeat_interleave(self.group_size, dim=1) ## (B,kv_heads,S,d_k) --> (B,h,S,d_k)
        value = value.repeat_interleave(self.group_size, dim=1) ## (B,kv_heads,S,d_k) --> (B,h,S,d_k)
 
 
        x, _ = GroupedQueryAttentionBlock.attention(query, key, value, mask, self.dropout)
        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.d_model) ## (B,h,S,d_k) ---> (B,S,h,d_k) --> (B,S,d)
        
        return self.w_o(x) ## (B,S,d) --> (B,S,d)
 
    @staticmethod
    def attention(query, key, value, mask, dropout: nn.Dropout):
        d_k = query.shape[-1]
        
        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k) ## (B,h,S,d_k) * (B,h,d_k,S) --> (B,h,S,S)
        if mask is not None:
            attention_scores.masked_fill_(mask == 0, torch.finfo(attention_scores.dtype).min)
        
        attention_scores = attention_scores.softmax(dim=-1) ## (B,h,S,S)
        
        if dropout is not None:
            attention_scores = dropout(attention_scores)
        
        return (attention_scores @ value), attention_scores ## (B,h,S,S) * (B,h,S,d_k) --> (B,h,S,d_k)

Gains:

Quality close to MHA
Comparable speed to MQA

Drawbacks:

More complex than MHA or MQA
How to find the optimal number of Groups - G?
Assignment of KV to right group

Multi Head Latent Attention (MLA)

Multi_Head_Latent_Attention Instead of reducing key-value heads like GQA, MLA compresses the key-value information into a much lower dimensional latent representation and caches that latent representation.

The hidden representation at time $t$ is $h_{t} \in R^{d}$ , MLA first applies a down projection:

c_{t}^{K V} = W_{D}^{K V} h_{t}, c_{t}^{K V} \in R^{d_{c}} (d_{c} << d)

So we get $c_{t}^{K V}$ which is a compressed latent representation containing the information needed for $K & V$ that is cached.

To recover keys and values during inference:

k_{t} = W_{U}^{K} c_{t}^{K V} v_{t} = W_{U}^{V} c_{t}^{K V}

The model learns a shared low-dimensional latent representation from which the $K & V$ representations can be generated.

The same idea is applied to compress queries $Q$ :

c_{t}^{Q} = W_{D}^{Q} h_{t}

and recovered as:

q_{t} = W_{U}^{Q} c_{t}^{Q}

$K V$ compressions is motivated by inference-time KV-cache reduction whereas $Q$ compression is additionally useful for reducing training memory.

[TODO]

Abdul Kadir's Virtual Home 🐈

Explorer

Attention

Multi Head Attention (MHA)

Longformer Attention

Sliding-Window Attention

Dilating Sliding-Window Attention

Global Attention

Multi Query Attention (MQA)

Grouped Query Attention (GQA)

Multi Head Latent Attention (MLA)

Graph View

Table of Contents

Backlinks