Perceiver Resampler Forward Pass

1. Background Knowledge

The Perceiver Resampler is a critical component in Vision-Language Models (VLMs) like Flamingo, designed to bridge the gap between high-dimensional, variable-length visual features and the fixed-length token sequences expected by language models. Unlike standard cross-attention where queries attend only to keys/values from a different source, the Perceiver Resampler introduces a unique mechanism: latent tokens that serve as both queries and part of the key/value sequence. This allows the latents to not only extract information from the image features but also to mix information among themselves, creating a more coherent and context-aware representation.

The core operation is a modified multi-head attention mechanism. In standard attention, $Q$, $K$, and $V$ are derived from separate inputs. Here, the queries $Q$ are derived solely from the latent tokens $L$. However, the keys $K$ and values $V$ are derived from the concatenation of the image features $X_f$ and the latent tokens $L$. This means each latent token attends to all image features and all other latent tokens. The mathematical formulation is: $Q = L W_Q, \quad K = [X_f; L] W_K, \quad V = [X_f; L] W_V$ where $[X_f; L]$ denotes vertical concatenation. This structure enables the model to refine the latent representations iteratively, making them robust summaries of the visual input.

Understanding scaled dot-product attention is essential. For each head, the attention weights are computed as $\text{softmax}(\frac{Q K^T}{\sqrt{d_h}})$, where $d_h$ is the dimension of each head. The scaling factor $\frac{1}{\sqrt{d_h}}$ prevents the dot products from growing too large, which would push the softmax function into regions with extremely small gradients. Finally, the outputs from all heads are concatenated and projected via $W_O$, with a residual connection adding the original latents $L$ to the result, ensuring stable gradient flow.

2. Algorithm Approach

The problem requires implementing a single forward pass of a Perceiver Resampler block using nested lists instead of tensor libraries like PyTorch. The approach involves manually simulating matrix operations and attention mechanisms.

1. Matrix Multiplication: Implement a helper function for matrix multiplication ($A @ B$) since standard libraries are not available.
2. Projection: Compute $Q$, $K_{feats}$, $V_{feats}$, $K_{latents}$, and $V_{latents}$ using the provided weight matrices $W_Q, W_K, W_V$.
3. Concatenation: Vertically concatenate the feature keys/values with the latent keys/values to form the full $K$ and $V$ matrices.
4. Multi-Head Attention:

• Split the projected matrices into num_heads by reshaping the last dimension.
• For each head, compute the scaled dot-product attention: calculate scores, apply softmax, and multiply by values.
• Concatenate the head outputs back into a single matrix.

5. Output Projection & Residual: Multiply the concatenated attention output by $W_O$ and add the original latents matrix.
6. Rounding: Round all final values to 4 decimal places.

3. Step-by-Step Strategy

1. Helper Functions:

• Create matmul(A, B) to perform matrix multiplication on nested lists.
• Create transpose(A) to transpose a matrix.
• Create softmax(row) to apply softmax to a 1D list, handling numerical stability by subtracting the max value.

2. Compute Projections:

• Calculate $Q = \text{matmul}(\text{latents}, W_Q)$.
• Calculate $K_{feats} = \text{matmul}(\text{feats}, W_K)$ and $V_{feats} = \text{matmul}(\text{feats}, W_V)$.
• Calculate $K_{latents} = \text{matmul}(\text{latents}, W_K)$ and $V_{latents} = \text{matmul}(\text{latents}, W_V)$.

3. Concatenate Keys and Values:

• Form $K = K_{feats} + K_{latents}$ (list concatenation of rows).
• Form $V = V_{feats} + V_{latents}$.
• Note: $Q$ remains separate with shape (n_latents, d).

4. Multi-Head Attention Loop:

• Determine head dimension $d_h = d / \text{num\_heads}$.
• Initialize an empty list for the final attention output.
• For each head $i$ from $0$ to $\text{num\_heads}-1$:
• Extract head-specific slices from $Q$, $K$, and $V$. For a matrix $M$, head $i$ corresponds to columns $i \cdot d_h$ to $(i+1) \cdot d_h$.
• Compute scores: $S = \text{matmul}(Q_{head}, \text{transpose}(K_{head}))$.
• Scale scores: Multiply each element in $S$ by $1/\sqrt{d_h}$.
• Apply softmax row-wise to $S$ to get attention weights $A$.
• Compute head output: $O_{head} = \text{matmul}(A, V_{head})$.
• Store $O_{head}$.

5. Concatenate Heads:

• Merge the outputs from all heads horizontally. If head $i$ output is $O_i$, the combined output for row $j$ is $[O_0[j], O_1[j], \dots]$.

6. Final Projection and Residual:

• Compute $\text{out} = \text{matmul}(\text{combined\_heads}, W_O)$.
• Add the original latents to out element-wise.
• Round every element to 4 decimal places.

4. Common Pitfalls

• Incorrect Concatenation: Ensure keys and values are concatenated vertically (adding rows), not horizontally. The sequence length becomes $n_{feats} + n_{latents}$.
• Head Splitting Logic: When splitting into heads, ensure you are slicing the columns of the projected matrices, not the rows. Each head processes a subset of the feature dimensions.
• Scaling Factor: Remember to scale by $1/\sqrt{d_h}$, where $d_h$ is the head dimension, not the total model dimension $d$.
• Softmax Stability: Always subtract the maximum value in each row before applying the exponential function in softmax to prevent overflow.
• Residual Connection: Do not forget to add the original latents to the final output. The problem specifies a residual connection.
• Rounding: Apply rounding only at the very end, after all computations are complete, to avoid cumulative precision errors.

5. Time & Space Complexity

• Time Complexity:

• Matrix multiplications dominate the cost. Let $N_L$ be $n_{latents}$, $N_F$ be $n_{feats}$, $D$ be the dimension, and $H$ be the number of heads.

• Projections: $O(N_L D^2 + N_F D^2)$.

• Attention Scores: $O(H \cdot N_L \cdot (N_F + N_L) \cdot D/H) = O(N_L (N_F + N_L) D)$.

• Final Projection: $O(N_L D^2)$.

• Overall: $O(D^2 (N_L + N_F) + N_L (N_F + N_L) D)$.

• Space Complexity:

• Storing intermediate matrices $Q, K, V$ and their head splits requires $O(D^2 + N_F D + N_L D)$ space.

• Attention score matrix size is $O(N_L (N_F + N_L))$.

• Overall: $O(D^2 + N_F D + N_L D + N_L (N_F + N_L))$.