Multi-Head Attention Forward Pass (Simplified)

hardPROTransformersNeural Networks, Matrix MathAdvanced 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 = [[1, 0, 1, 0, 1, 0], [0, 1, 0, 1, 0, 1]]

Output

[[2.0, 0.0], [0.0, 2.0]]

Reasoning

Inputs: X is $N = 2$ tokens with $D = 2$ features, the identity $I$ . W_QKV is $2 \times 6$ : three $2 \times 2$ blocks (query, key, value), each $I$ , so $W_{QKV} = [\,I \mid I \mid I\,]$ .
Projection: $X W_{QKV}$ is $2 \times 6$ ; its column blocks of width $D$ are $Q$ , $K$ , $V$ . With identity blocks, $Q = K = V = X = I$ .
Scores: $A = Q K^T$ is $N \times N$ ; entry $a_{ij}$ is the dot product of token $i$ 's query with token $j$ 's key. Here $a_{11} = a_{22} = 1$ and $a_{12} = a_{21} = 0$ , so $A = I$ : each token attends only to itself, with raw weight 1 (no scaling, no softmax).
Mix + residual: $A V = I$ , so $Y = X + A V = I + I$ , giving [[2.0, 0.0], [0.0, 2.0]], an $N \times D$ result like $X$ .
Notice: with all-ones $X$ , every score is 2, $AV$ is all 4s and $Y$ is all 5s: unnormalised scores grow fast.

Constraints

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

Learn this in

Transformer ArchitectureLLM Study Plan NLP FundamentalsML Study Plan Attention MechanismsNLP Study Plan

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

Concept primerPREVIEW

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.

PRO

Continue the full explanation

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

PRO

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