Multi-Head Cross-Attention Output

mediumPROBridge Modules: Projectors, Q-Former & Cross-AttentionVision-Language ModelsVLM 1: Vision-Language Models

Problem Statement

Extend single-head cross-attention to h heads: split the model dimension across heads, run scaled dot-product attention independently per head, then concatenate. This is the resampler/Q-Former core.

Background

With model dimension d_model and h heads, each head works in dimension d_k = d_model / h. Given already-projected Q (n_q x d_model), K, V (n_kv x d_model), split each along the feature axis into h contiguous blocks. For head j:

$O_j = \text{softmax}\!\left(\frac{Q_j K_j^\top}{\sqrt{d_k}}\right) V_j$

Concatenate O_0, ..., O_{h-1} back to n_q x d_model. (No output projection here — just the attention.) Softmax is row-wise and numerically stable.

Your Task

Implement:

def multihead_cross_attention(Q, K, V, h):

Return the n_q x d_model output as a nested list rounded to 4 decimals.

Input Format

Q: n_q x d_model; K, V: n_kv x d_model.
h (int): number of heads; d_model divisible by h.

Output Format

An n_q x d_model nested list rounded to 4 decimals.

Sample

Q = [[1.0, 0.0]]
K = [[1.0, 0.0], [0.0, 1.0]]
V = [[2.0, 0.0], [0.0, 3.0]]
print(multihead_cross_attention(Q, K, V, 2))

Output:

[[1.4621, 1.5]]

Example

Input

Q = [[1.0, 0.0]]
K = [[1.0, 0.0], [0.0, 1.0]]
V = [[2.0, 0.0], [0.0, 3.0]]
print(multihead_cross_attention(Q, K, V, 2))

Output

[[1.4621, 1.5]]

Reasoning

Two heads of size 1 (d_k=1, scale=1). Head 0: Q=[1], K=[[1],[0]] gives scores [1,0], softmax [0.7311,0.2689], attending V0=[2],V1=[0] -> 1.4621. Head 1: Q=[0], K=[[0],[1]] gives scores [0,1], softmax [0.2689,0.7311], attending V0=[0],V1=[3] -> 1.5. Concatenated: [1.4621, 1.5].

Constraints

d_model % h == 0; d_k = d_model / h.
Each head scales by 1/sqrt(d_k) and softmaxes over the key axis, stably.
Concatenate heads in order; round to 4 decimals; avoid -0.0.

Learn this in

Latent Diffusion & Stable DiffusionDiffusion Study Plan Multimodal LLMsLLM Study Plan Transformer ArchitectureLLM Study Plan Attention MechanismsNLP Study Plan

More from VLM 1: Vision-Language Models

Visual Token Pruning by Attention Masshard CHAIR-i and CHAIR-s Hallucination Ratesmedium Merge a LoRA Update into the Base Weightmedium QLoRA Memory Footprint Estimatemedium POPE Precision, Recall, and F1medium

Concept primerPREVIEW

1. Background Knowledge

Multi-head attention is the core mechanism in Transformers that allows a model to jointly attend to information from different representation subspaces. Instead of performing a single attention operation over the full model dimension $d_{\text{model}}$ , the input is split into $h$ heads, each operating in a lower-dimensional subspace of size $d_k = d_{\text{model}} / h$ . This enables the model to capture diverse relationships (e.g., syntactic vs. semantic) in parallel.

In cross-attention, the query matrix $Q$ comes from one sequence (e.g., text tokens or learnable queries in a Q-Former), while the key $K$ and value $V$ matrices come from a different sequence (e.g., vision features). The attention weights are computed as $\text{softmax}(QK^\top / \sqrt{d_k})$ , and the output is a weighted sum of the values. The scaling factor $\sqrt{d_k}$ prevents the dot products from growing too large, which would push the softmax into regions with vanishing gradients.

In vision-language models like BLIP-2, the Q-Former uses multi-head cross-attention to bridge visual and textual modalities. Learnable query tokens attend to visual features, effectively "resampling" the visual information into a fixed-size representation that the language model can consume. This problem isolates that cross-attention computation without the surrounding projection layers.

2. Algorithm Approach

The algorithm follows a split → compute → concatenate pattern:

Split: Partition each of $Q$ , $K$ , $V$ along the feature (column) axis into $h$ contiguous blocks of width $d_k$ .
Per-head attention: For each head $j$ , compute the scaled dot-product attention: $O_j = \text{softmax}(Q_j K_j^\top / \sqrt{d_k}) \, V_j$ .
Concatenate: Horizontally stack all $O_j$ blocks to reconstruct the full $n_q \times d_{\text{model}}$ output.

This is a direct generalization of single-head attention. The key insight is that each head operates independently on its own slice of the feature space, and the results are simply concatenated—no mixing across heads occurs at this stage (that would be the role of an output projection, which is explicitly excluded here).

3. Step-by-Step Strategy

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