📘

Multi-Head Attention Forward Pass (Simplified)

HardNeural Networks, Matrix Math 📚 Advanced Architectures

Implement a simplified Multi-Head Self-Attention (MHSA) forward pass.

Given input embeddings $X \in \mathbb{R}^{N \times D}$ and weight matrix $W_{QKV} \in \mathbb{R}^{D \times 3D}$ :

Compute $[Q, K, V] = X \cdot W_{QKV}$ (split into three D-dimensional matrices)
Compute attention: $A = Q \cdot K^T$ (unscaled, no softmax)
Apply residual: $Y = X + A \cdot V$

Simplifications:

Single head (no multi-head concatenation)
No softmax or scaling factor $\sqrt{D}$
No output projection $W_O$
No layer normalization or MLP

Example:

Input:

X = [[1,0],[0,1]], W_QKV = 2x6 identity-like matrix

Output:

2x2 matrix Y after residual connection

Reasoning:

Q, K, V are derived from X @ W_QKV split into thirds. A = QK^T computes attention scores. Y = X + AV applies residual.

Constraints:

$N$ (Sequence Length): 2-4
$D$ (Embedding Dimension): 2-3
Input $X$ and $W_{QKV}$ are provided as numpy arrays

More from Advanced Architectures

Complete Mixer Block with ResidualsHard Channel Mixing with GELUMedium Token Mixing LayerMedium Bidirectional Scan for VisionMedium Discretize Continuous State Space ModelHard

1. Background Knowledge

Multi-Head Self-Attention (MHSA) is the core mechanism of Transformer models, enabling sequences to attend to different representation subspaces in parallel. In self-attention, input embeddings $X \in \mathbb{R}^{N \times D}$ (where $N$ is sequence length, $D$ is embedding dimension) are linearly projected into Query ( $Q$ ), Key ( $K$ ), and Value ( $V$ ) matrices via a shared weight matrix $W_{QKV} \in \mathbb{R}^{D \times 3D}$ .

The standard attention computes similarity scores as $A =$ \text{softmax} $\left($ \frac{QK^T}{ $\sqrt{d_k}$ } $\right)V$ , where each row of $A$ represents attention weights for a position across the sequence. This problem simplifies to unscaled dot-product attention ( $A = QK^T$ ) without softmax, followed by a residual connection $Y = X + AV$ . Multi-head variants split $Q,K,V$ into $h$ heads, compute attention independently, then concatenate and project—but here it's single-head for simplicity.

Prerequisites: Matrix multiplication, NumPy array slicing, understanding $X W_{QKV}$ yields concatenated $[Q|K|V]$ each of shape $\mathbb{R}^{N \times D}$ .

2. Algorithm Approach

The forward pass follows the standard Transformer encoder attention block, stripped to essentials:

Linear projection: $Z = X W_{QKV} \in \mathbb{R}^{N \times 3D}$
Split: $Q = Z[:, :D]$ , $K = Z[:, D:2D]$ , $V = Z[:, 2D:]$ (each $\mathbb{R}^{N \times D}$ )
Attention scores: $A = Q K^T \in \mathbb{R}^{N \times N}$ (raw dot products)
Weighted values: $AV \in \mathbb{R}^{N \times D}$
Residual: $Y = X + AV \in \mathbb{R}^{N \times D}$

This mirrors the mathematical foundation in Transformer papers, focusing on pairwise sequence interactions without scaling or normalization.

👑

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

🔒

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.