Scaled Dot-Product Cross-Attention

1. Background Knowledge

Cross-attention is a fundamental mechanism in Vision-Language Models (VLMs) that allows one sequence of tokens (queries) to attend to information from a different sequence (keys and values). Unlike self-attention, where queries, keys, and values are derived from the same input, cross-attention decouples these roles. In VLMs, the queries typically come from the language model or a learned set of resampler tokens, while the keys and values originate from the vision encoder’s patch features. This enables the language side to selectively "read" relevant visual information.

The core computation is defined as:

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

Here, $Q \in \mathbb{R}^{n_q \times d_k}$, $K \in \mathbb{R}^{n_{kv} \times d_k}$, and $V \in \mathbb{R}^{n_{kv} \times d_v}$. The output shape is $(n_q, d_v)$, meaning the number of output tokens matches the number of queries, not the keys. This is crucial for tasks like image captioning, where a fixed number of text tokens must aggregate information from a variable number of image patches.

Two critical numerical considerations arise:

1. Scaling by $\sqrt{d_k}$: Without this scaling, dot products between high-dimensional vectors tend to have large magnitudes, causing the softmax function to saturate (outputs approach 0 or 1). This leads to vanishing gradients during training. Dividing by $\sqrt{d_k}$ keeps the variance of the dot products stable.
2. Numerical Stability of Softmax: Directly computing $\exp(x)$ can overflow for large $x$. The standard trick is to subtract the maximum value in each row before exponentiating: $\text{softmax}(x_i) = \frac{\exp(x_i - \max(x))}{\sum_j \exp(x_j - \max(x))}$. This does not change the result but prevents overflow.

2. Algorithm Approach

The problem requires implementing the scaled dot-product cross-attention from scratch using nested lists. The approach involves three main phases:

1. Matrix Multiplication ($QK^\top$): Compute the dot product between each query vector and each key vector. This results in an attention score matrix of shape $(n_q, n_{kv})$.
2. Scaling and Softmax: Scale the scores by $1/\sqrt{d_k}$ and apply a numerically stable softmax row-wise to obtain attention weights. These weights sum to 1 for each query.
3. Weighted Sum ($AV$): Multiply the attention weight matrix by the value matrix $V$ to produce the final output. Each output token is a weighted sum of the value vectors, where the weights are determined by the attention scores.

Since the input is given as nested lists, you must implement matrix operations manually or use helper functions for dot products and matrix multiplication. Avoid using external libraries like NumPy unless explicitly allowed; the problem likely tests your understanding of the underlying mechanics.

3. Step-by-Step Strategy

1. Extract Dimensions: Determine $n_q$ (number of queries), $n_{kv}$ (number of keys/values), $d_k$ (dimension of keys/queries), and $d_v$ (dimension of values) from the input lists.
2. Compute Attention Scores:

• Initialize a matrix scores of shape $(n_q, n_{kv})$.
• For each query $i$ and key $j$, compute the dot product $Q[i] \cdot K[j]$.
• Store the result in scores[i][j].

3. Scale the Scores:

• Calculate the scaling factor $s = 1/\sqrt{d_k}$.
• Multiply every element in scores by $s$.

4. Apply Numerically Stable Softmax:

• For each row $i$ in scores:
• Find the maximum value $m_i = \max_j(\text{scores}[i][j])$.
• Subtract $m_i$ from each element in the row.
• Exponentiate each element: $\exp(\text{scores}[i][j] - m_i)$.
• Sum the exponentiated values to get the denominator $S_i$.
• Divide each exponentiated element by $S_i$ to get the attention weights $A[i][j]$.

5. Compute Output:

• Initialize an output matrix out of shape $(n_q, d_v)$.
• For each query $i$ and value dimension $k$:
• Compute the weighted sum: $\text{out}[i][k] = \sum_{j=0}^{n_{kv}-1} A[i][j] \cdot V[j][k]$.

6. Round and Return:

• Round each element in out to 4 decimal places.
• Return the nested list.

4. Common Pitfalls

• Incorrect Scaling: Dividing by $d_k$ instead of $\sqrt{d_k}$ will cause the attention scores to be too small, leading to nearly uniform attention weights and poor performance. Always use $\sqrt{d_k}$.
• Softmax Overflow: Failing to subtract the row maximum before exponentiating can lead to inf values if the scores are large. Always implement the stable softmax variant.
• Shape Mismatch: Ensure that the output shape is $(n_q, d_v)$, not $(n_{kv}, d_v)$. The number of output tokens is determined by the number of queries.
• Manual Matrix Multiplication Errors: When implementing dot products manually, ensure you are iterating over the correct dimensions. For $QK^\top$, you are dotting rows of $Q$ with rows of $K$. For $AV$, you are dotting rows of $A$ with columns of $V$.
• Rounding Precision: Round only at the final step. Intermediate rounding can accumulate errors and lead to incorrect final results.

5. Time & Space Complexity

• Time Complexity:

• Computing $QK^\top$: $O(n_q \cdot n_{kv} \cdot d_k)$.

• Softmax: $O(n_q \cdot n_{kv})$.

• Computing $AV$: $O(n_q \cdot n_{kv} \cdot d_v)$.

• Overall: $O(n_q \cdot n_{kv} \cdot (d_k + d_v))$. This is dominated by the matrix multiplications.

• Space Complexity:

• Attention Scores Matrix: $O(n_q \cdot n_{kv})$.

• Attention Weights Matrix: $O(n_q \cdot n_{kv})$.

• Output Matrix: $O(n_q \cdot d_v)$.

• Overall: $O(n_q \cdot n_{kv} + n_q \cdot d_v)$. If we reuse the scores matrix for weights, it can be reduced to $O(n_q \cdot n_{kv})$.