Scaled Dot-Product Cross-Attention

Problem Statement

Compute the output of a single-head cross-attention layer given explicit query, key and value matrices - the operation that lets learned queries (or text tokens) read from image features.

Background

Cross-attention is the same computation as self-attention except that the queries come from one stream and the keys/values from another. In a VLM the queries are language-side (or the Q-Former's learned queries) and the keys/values are the vision encoder's patch features:

$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$

with Q of shape (n_q, d_k), K of shape (n_kv, d_k), V of shape (n_kv, d_v), and the output of shape (n_q, d_v). The output length follows the queries, not the keys - which is exactly why a resampler with 64 queries turns 576 patch features into 64 tokens.

Two things trip people up:

• The 1/sqrt(d_k) scale is not cosmetic. Dot products of d_k independent terms grow like sqrt(d_k); without the divisor the softmax saturates and gradients vanish. Divide by sqrt(d_k), never by d_k.
• Softmax is row-wise over the keys and must be numerically stable: subtract each row's maximum before exponentiating.

Your Task

Implement:

def cross_attention(Q, K, V):

Return the (n_q, d_v) output as a nested list with every entry rounded to 4 decimals.

Input Format

• Q - nested list, shape (n_q, d_k)
• K - nested list, shape (n_kv, d_k)
• V - nested list, shape (n_kv, d_v)

Output Format

A nested list of n_q rows of d_v floats, each rounded to 4 decimals.

Sample

Q = [[1.0, 0.0]]
K = [[1.0, 0.0], [0.0, 1.0]]
V = [[1.0, 2.0], [3.0, 4.0]]
print(cross_attention(Q, K, V))

Output:

[[1.6605, 2.6605]]