Explain the Self-Attention mechanism in Transformer architectures and how its computational complexity scales with sequence length.

Quick Revision Summary

Transformer Self-Attention Mechanism Self-attention allows an input token to dynamically attend to every other token in the sequence, computing context-aware contextual represen...

Interactive Practice Mode

Formulate your thoughts or take notes before revealing the model solution.

Model Answer & Deep-Dive Explanation

Transformer Self-Attention Mechanism

Self-attention allows an input token to dynamically attend to every other token in the sequence, computing context-aware contextual representations.

Mathematical Formulation

Given an input sequence matrix (X in mathbb{R}^{n imes d}), three projection matrices produce Queries ((Q)), Keys ((K)), and Values ((V)):

[ Q = X W_Q, quad K = X W_K, quad V = X W_V ]

The scaled dot-product attention is defined as:

[ ext{Attention}(Q, K, V) = ext{softmax}left( rac{Q K^T}{sqrt{d_k}} ight) V ]

Where (d_k) is the key dimension, preventing dot products from growing excessively large and pushing the softmax function into regions with near-zero gradients.

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V):
    d_k = Q.size(-1)
    # Q: [batch, heads, seq_len, d_k]
    # K: [batch, heads, seq_len, d_k]
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    attention_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attention_weights, V)
    retu
 output, attention_weights

Computational Complexity

1. Matrix Multiplication (Q K^T): (mathcal{O}(n^2 cdot d)) operations.

Softmax over (n imes n): (mathcal{O}(n^2)) time.

Multiplication with (V): (mathcal{O}(n^2 cdot d)) time.

Total Time Complexity: (mathcal{O}(n^2 cdot d))
Total Space Complexity: (mathcal{O}(n^2))

Interviewer Evaluation Points & Scoring Rubric

1 Explain trade-offs between memory, CPU, and network latency.
2 Identify failure modes, boundary conditions, and scalability bottlenecks.
3 Discuss caching strategies, data consistency models, and recovery.