08. Multi-Head Attention

flowchart TD
    INPUT["'The cat sat on the mat'"]

    INPUT --> H1["Head 1\n(Syntactic Role)\n→ Connects verbs to subjects\n'sat' ↔ 'cat'"]

    INPUT --> H2["Head 2\n(Coreference)\n→ Connects pronouns to nouns\n'it' ↔ 'animal'"]

    INPUT --> H3["Head 3\n(Semantic Relatedness)\n→ Groups related concepts\n'cat' ↔ 'mat' (cats sit on mats)"]

    INPUT --> H4["Head 4\n(Positional Proximity)\n→ Connects nearby words\n'sat' ↔ 'on'"]

    H1 & H2 & H3 & H4 --> COMBINE["Combine all views\n(Richer understanding)"]
    COMBINE --> OUT["Final: Each word understands\nmultiple relationships at once"]

    style H1 fill:#3b82f6,color:#fff
    style H2 fill:#22c55e,color:#fff
    style H3 fill:#f59e0b,color:#fff
    style H4 fill:#8b5cf6,color:#fff
    style COMBINE fill:#ef4444,color:#fff
    style OUT fill:#22c55e,color:#fff

flowchart TD
    REPORT["Same Report 📊"]

    REPORT --> FIN["Finance Director\n▶ Focuses on numbers"]
    REPORT --> HR["HR Director\n▶ Focuses on people"]
    REPORT --> MKT["Marketing Director\n▶ Focuses on perception"]
    REPORT --> OPS["Operations Director\n▶ Focuses on processes"]

    FIN & HR & MKT & OPS --> MEETING["Board Meeting\n▶ Combine insights"]
    MEETING --> DECISION["Better understanding\nthan any one director\ncould achieve alone"]

    style FIN fill:#3b82f6,color:#fff
    style HR fill:#22c55e,color:#fff
    style MKT fill:#f59e0b,color:#fff
    style OPS fill:#8b5cf6,color:#fff
    style MEETING fill:#ef4444,color:#fff
    style DECISION fill:#22c55e,color:#fff

flowchart TD
    ANIMAL["'animal'"]

    subgraph ONE_HEAD["Single Head (Compromised)"]
        O1["'didn't cross': 0.40"]
        O2["'it': 0.25"]
        O3["'tired': 0.20"]
        O4["'the': 0.05"]
        O5["'street': 0.05"]
        O6["'because': 0.03"]
        O7["'was': 0.02"]
        NOTE1["One weight per pair.\nMust distribute across all needs."]
    end

    subgraph MULTI_HEAD["Multi-Head (Specialized)"]
        H1_HEAD["Head 1 (Syntax):\n'didn't cross': 0.90\n'it': 0.03\nOther: low"]
        H2_HEAD["Head 2 (Coref):\n'it': 0.85\n'didn't cross': 0.05\nOther: low"]
        H3_HEAD["Head 3 (Causal):\n'tired': 0.80\n'didn't cross': 0.10\nOther: low"]
    end

    style ANIMAL fill:#22c55e,color:#fff
    style ONE_HEAD fill:#ef4444,color:#fff
    style MULTI_HEAD fill:#22c55e,color:#fff

flowchart TD
    SCENE["Same Evidence\n(Crime Scene)"]

    SCENE --> DA["Detective A\n(Forensics)\n→ Fingerprints"]
    SCENE --> DB["Detective B\n(Finance)\n→ Bank records"]
    SCENE --> DC["Detective C\n(Psychology)\n→ Motives"]
    SCENE --> DD["Detective D\n(Timeline)\n→ Alibis"]

    DA & DB & DC & DD --> WB["Whiteboard\n(Combine findings)"]
    WB --> SOLVE["Richer understanding\nthan one detective alone"]

    style DA fill:#3b82f6,color:#fff
    style DB fill:#22c55e,color:#fff
    style DC fill:#f59e0b,color:#fff
    style DD fill:#8b5cf6,color:#fff
    style WB fill:#ef4444,color:#fff
    style SOLVE fill:#22c55e,color:#fff

Head 1:    W_Q¹, W_K¹, W_V¹
Head 2:    W_Q², W_K², W_V²
Head 3:    W_Q³, W_K³, W_V³
...
Head 8:    W_Q⁸, W_K⁸, W_V⁸ (original Transformer)
Head 96:   W_Q⁹⁶, W_K⁹⁶, W_V⁹⁶ (GPT-3)

Head 1: Q¹ = X × W_Q¹, K¹ = X × W_K¹, V¹ = X × W_V¹
       Attention¹ = softmax(Q¹ × K¹ᵀ / √d_k) × V¹

Head 2: Q² = X × W_Q², K² = X × W_K², V² = X × W_V²
       Attention² = softmax(Q² × K²ᵀ / √d_k) × V²

[Head 1 output, Head 2 output, ..., Head H output]

Final = Concat[Head 1, ..., Head H] × W_O

flowchart TD
    INPUT["Input Vectors X"]

    INPUT --> H1["Head 1\nQ¹=W_Q¹×X, K¹=W_K¹×X, V¹=W_V¹×X\n→ Self-Attention"]
    INPUT --> H2["Head 2\nQ²=W_Q²×X, K²=W_K²×X, V²=W_V²×X\n→ Self-Attention"]
    INPUT --> H3["Head 3\n..."]
    INPUT --> HN["Head H\n..."]

    H1 --> OUT1["Output 1\n(dim = d_k)"]
    H2 --> OUT2["Output 2\n(dim = d_k)"]
    H3 --> OUT3["Output 3\n(dim = d_k)"]
    HN --> OUTN["Output H\n(dim = d_k)"]

    OUT1 & OUT2 & OUT3 & OUTN --> CONCAT["Concatenate\n(dim = d_k × H)"]
    CONCAT --> PROJ["Linear Projection W_O\n(dim → d_model)"]
    PROJ --> FINAL["Final Output\n(Each token now has\nmulti-dimensional understanding)"]

    style INPUT fill:#3b82f6,color:#fff
    style H1 fill:#3b82f6,color:#fff
    style H2 fill:#22c55e,color:#fff
    style H3 fill:#f59e0b,color:#fff
    style HN fill:#8b5cf6,color:#fff
    style CONCAT fill:#ef4444,color:#fff
    style PROJ fill:#8b5cf6,color:#fff
    style FINAL fill:#22c55e,color:#fff

flowchart TD
    BERT["BERT-base (12 layers × 12 heads = 144 total)"]

    BERT --> L1["Layer 1: Low-level patterns"]
    L1 --> L1H1["Head 1: Next word prediction\n(the → cat, sat → on)"]
    L1 --> L1H2["Head 2: Position patterns\n(token i attends to token i+1)"]

    BERT --> L5["Layer 5: Syntactic patterns"]
    L5 --> L5H1["Head 1: Subject-verb\n(cat → sat)"]
    L5 --> L5H2["Head 3: Adjective-noun\n(red → car, big → house)"]

    BERT --> L10["Layer 10: Semantic patterns"]
    L10 --> L10H1["Head 2: Coreference\n(it → animal, she → doctor)"]
    L10 --> L10H2["Head 4: Same entity\n(Tesla → company, they → team)"]

    BERT --> L12["Layer 12: Sentence-level"]
    L12 --> L12H1["Head 5: [CLS] → all tokens\n(sentence representation)"]

    style L1 fill:#3b82f6,color:#fff
    style L5 fill:#22c55e,color:#fff
    style L10 fill:#f59e0b,color:#fff
    style L12 fill:#8b5cf6,color:#fff

import numpy as np

class MultiHeadAttention:
    def __init__(self, d_model=64, num_heads=8):
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # dimension per head
        self.d_model = d_model

        # Initialize QKV weight matrices for ALL heads at once
        # (efficient: one large matrix instead of H separate small ones)
        np.random.seed(42)
        self.W_Q = np.random.randn(d_model, d_model) * 0.1
        self.W_K = np.random.randn(d_model, d_model) * 0.1
        self.W_V = np.random.randn(d_model, d_model) * 0.1
        self.W_O = np.random.randn(d_model, d_model) * 0.1

    def split_heads(self, x):
        """Split the last dimension into (num_heads, d_k)."""
        batch_size, seq_len, _ = x.shape
        x = x.reshape(batch_size, seq_len, self.num_heads, self.d_k)
        return x.transpose(0, 2, 1, 3)  # (batch, heads, seq, d_k)

    def combine_heads(self, x):
        """Inverse of split_heads."""
        batch_size, heads, seq_len, d_k = x.shape
        x = x.transpose(0, 2, 1, 3)  # (batch, seq, heads, d_k)
        return x.reshape(batch_size, seq_len, self.d_model)

    def scaled_dot_product_attention(self, Q, K, V):
        """Compute attention for a single batch of heads."""
        scores = np.matmul(Q, K.transpose(0, 1, 3, 2))
        scores = scores / np.sqrt(self.d_k)

        # Softmax
        exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
        attention = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)

        return np.matmul(attention, V), attention

    def forward(self, X):
        """X shape: (batch_size, seq_len, d_model)"""
        batch_size = X.shape[0]

        # 1. Linear projections for Q, K, V
        Q = np.dot(X, self.W_Q)  # (batch, seq, d_model)
        K = np.dot(X, self.W_K)
        V = np.dot(X, self.W_V)

        # 2. Split into multiple heads
        Q = self.split_heads(Q)  # (batch, heads, seq, d_k)
        K = self.split_heads(K)
        V = self.split_heads(V)

        # 3. Apply attention to each head independently
        attn_output, attention = self.scaled_dot_product_attention(Q, K, V)

        # 4. Combine heads
        attn_output = self.combine_heads(attn_output)

        # 5. Final projection
        output = np.dot(attn_output, self.W_O)

        return output, attention


# Example usage
d_model = 64
batch_size = 1
seq_len = 4  # "The cat sat"

X = np.random.randn(batch_size, seq_len, d_model)

mha = MultiHeadAttention(d_model=d_model, num_heads=8)
output, attention = mha.forward(X)

print(f"Input shape: {X.shape}")
print(f"Output shape: {output.shape}")
print(f"Attention shape: {attention.shape}")
# Input:  (1, 4, 64)   — 4 tokens, 64-dim vectors
# Output: (1, 4, 64)   — same shape (content is now richer)
# Attention: (1, 8, 4, 4) — 8 heads, each with 4×4 attention matrix

# Show attention for first head, first token
print(f"\nHead 0, Token 0 attention weights:")
print(np.round(attention[0, 0, 0], 2))

class MultiHeadAttention {
  constructor(dModel, numHeads) {
    this.numHeads = numHeads;
    this.dK = dModel / numHeads;
    this.dModel = dModel;

    const rand = () => (Math.random() - 0.5) * 0.2;

    // Create weight matrices for each head (simplified)
    this.heads = Array.from({ length: numHeads }, () => ({
      W_Q: Array.from({ length: dModel }, () =>
        Array.from({ length: this.dK }, () => rand())),
      W_K: Array.from({ length: dModel }, () =>
        Array.from({ length: this.dK }, () => rand())),
      W_V: Array.from({ length: dModel }, () =>
        Array.from({ length: this.dK }, () => rand())),
    }));

    // Output projection
    this.W_O = Array.from({ length: dModel }, () =>
      Array.from({ length: dModel }, () => rand()));
  }

  forward(X) {
    // Run each head independently
    const headOutputs = this.heads.map(head => {
      const { Q, K, V } = this._computeQKV(X, head);
      return this._attention(Q, K, V);
    });

    // Concatenate all head outputs
    const concat = X.map((tokenVec, i) =>
      headOutputs.flatMap(headOut => headOut[i])
    );

    // Output projection
    return concat.map(vec =>
      this.W_O[0].map((_, j) =>
        vec.reduce((sum, v, k) => sum + v * this.W_O[k][j], 0)
      )
    );
  }

  _computeQKV(X, head) {
    const matMul = (mat, vec) =>
      mat[0].map((_, j) => mat.reduce((s, r, i) => s + r[j] * vec[i], 0));

    return {
      Q: X.map(v => matMul(head.W_Q, v)),
      K: X.map(v => matMul(head.W_K, v)),
      V: X.map(v => matMul(head.W_V, v)),
    };
  }

  _attention(Q, K, V) {
    const n = Q.length;
    const scores = Q.map(q =>
      K.map(k => q.reduce((sum, v, i) => sum + v * k[i], 0))
    );

    // Scale
    const scaled = scores.map(row =>
      row.map(s => s / Math.sqrt(this.dK))
    );

    // Softmax
    const softmax = arr => {
      const max = Math.max(...arr);
      const exp = arr.map(x => Math.exp(x - max));
      const sum = exp.reduce((a, b) => a + b, 0);
      return exp.map(x => x / sum);
    };

    const attn = scaled.map(row => softmax(row));

    // Weighted values
    return V[0].map((_, j) =>
      attn.map((row, i) =>
        row.reduce((sum, w, k) => sum + w * V[k][j], 0)
      )
    );
  }
}

// Test
const mha = new MultiHeadAttention(64, 8);
const input = Array.from({ length: 4 }, () =>
  Array.from({ length: 64 }, () => Math.random() - 0.5));

const output = mha.forward(input);
console.log('Input tokens:', input.length);
console.log('Output tokens:', output.length);
console.log('Output dim per token:', output[0].length);