Self-Attention Mechanism

Self-attention is a fundamental component of transformer architectures that allows models to weigh the importance of different parts of the input sequence when processing each element. The basic self-attention formula is:

Where:

• Q (Query), K (Key), and V (Value) are learned linear transformations of the input
• \(d_k\) is the dimension of the key vectors
• The scaling factor \(\sqrt{d_k}\) prevents the dot products from growing too large

Encoder vs Decoder Attention

In the encoder:

• Uses bidirectional self-attention where each position can attend to all other positions
• The attention mask is all ones, allowing full visibility
• Mathematically represented as: \(M_{enc} = \mathbf{1}_{n \times n}\)

In the decoder:

• Uses masked self-attention to prevent positions from attending to subsequent positions
• The attention mask is lower triangular: \(M_{dec}[i,j] = \begin{cases} 0 & \text{if } j > i \\ 1 & \text{otherwise} \end{cases}\)
• Cross-attention to encoder outputs:

Parameter-Efficient Fine-Tuning (PEFT) with LoRA

Low-Rank Adaptation (LoRA) is a PEFT technique that modifies the attention weights efficiently:

Where:

• \(W_0\) is the frozen pretrained weights
• \(B \in \mathbb{R}^{d\times r}\) and \(A \in \mathbb{R}^{r\times d}\) are low-rank matrices
• \(r\) is the rank (typically 8 or 16)

Implementation Details

# Original attention projection
y = Wx

# LoRA equivalent
y = W_0x + BA(x)

# Efficient implementation
def lora_layer(x, W_0, A, B, alpha):
    return W_0(x) + alpha * B(A(x))
        

Understanding KV Cache

Key-Value (KV) caching is a crucial optimization technique used in transformer models during inference, particularly in autoregressive generation. It significantly reduces computational overhead by storing and reusing previously computed key and value tensors.

Why KV Cache Matters

• Eliminates redundant computations in autoregressive generation
• Reduces memory bandwidth requirements
• Speeds up inference by up to 4x in typical scenarios

Implementation in Decoder Architecture

class DecoderWithKVCache:
    def __init__(self):
        self.kv_cache = {}
    
    def forward(self, x, cache_position):
        # Generate new key and value
        new_k = self.k_proj(x)  # [B, 1, H]
        new_v = self.v_proj(x)  # [B, 1, H]
        
        # Update cache
        if cache_position in self.kv_cache:
            k = torch.cat([self.kv_cache[cache_position][0], new_k], dim=1)
            v = torch.cat([self.kv_cache[cache_position][1], new_v], dim=1)
        else:
            k, v = new_k, new_v
            
        self.kv_cache[cache_position] = (k, v)
        
        # Compute attention using cached keys and values
        q = self.q_proj(x)  # [B, 1, H]
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attn = F.softmax(attn, dim=-1)
        
        return torch.matmul(attn, v)

Memory Efficiency

The memory requirements for KV cache can be expressed as:

Where:

• \(N_{layers}\) is the number of transformer layers
• \(S_{seq}\) is the sequence length
• \(D_{model}\) is the model dimension
• \(B_{size}\) is the batch size

Performance Impact

Consider a typical inference scenario:

• Without KV Cache: \(O(n^2)\) complexity per token
• With KV Cache: \(O(n)\) complexity per token
• Practical speedup: 2-4x faster inference