One Perceiver Resampler Block
Problem Statement
Implement one full Perceiver-Resampler cross-attention block: learned queries attend to the concatenation of the image features and the queries themselves (the resampler's key/value trick), followed by an add-and-project residual.
Background
The Flamingo/BLIP-2 resampler updates Q learned latents by cross-attending to visual features X. A single block does:
- Keys/values come from KV = concat(X, latents) along the token axis — the latents attend to the image and to themselves.
- Single-head scaled dot-product attention with the latents as queries: A = softmax(latents @ KV^T / sqrt(d)) @ KV.
- Residual add: out = latents + A.
All matrices share dimension d (no separate projections in this simplified block). Softmax is row-wise and numerically stable.
Your Task
Implement:
def perceiver_block(latents, X):
- latents: Q x d learned queries.
- X: N x d image features.
Return the Q x d output as a nested list rounded to 4 decimals.
Input Format
- latents: Q x d nested list.
- X: N x d nested list.
Output Format
- A Q x d nested list rounded to 4 decimals.
Sample
latents = [[1.0, 0.0]]
X = [[0.0, 2.0]]
print(perceiver_block(latents, X))
Output:
[[1.6698, 0.6605]]
Example:
latents = [[1.0, 0.0]] X = [[0.0, 2.0]] print(perceiver_block(latents, X))
[[1.6698, 0.6605]]
KV = [[0,2],[1,0]]. Scores = latents·KV^T / sqrt(2) = [0, 1]/sqrt(2) = [0, 0.7071]; softmax = [0.3302, 0.6698]; A = 0.3302*[0,2]+0.6698*[1,0] = [0.6698, 0.6605]; out = latents + A = [1.6698, 0.6605].
Constraints:
1 <= Q, N <= 256,1 <= d <= 512.- Keys/values are
concat(X, latents)along the token axis. - Scale by
1/sqrt(d); softmax over the KV axis, stably; then residual-add the latents. - Round to 4 decimals; avoid
-0.0.
1. Background Knowledge
The Perceiver-Resampler is a bridge module used in vision-language models (e.g., Flamingo, BLIP-2) to compress a variable-length sequence of visual features into a fixed number of learned latent vectors. Unlike a standard cross-attention layer where queries come from one modality and keys/values from another, the resampler introduces a self-attention trick: the learned latents attend not only to the image features but also to themselves. This allows the latents to refine their own representations iteratively while still grounding them in the visual input.
The core operation is scaled dot-product attention. Given a query matrix Q, a key matrix K, and a value matrix V, the attention output is computed as:
Attention(Q,K,V)=softmax(dQKT)Vwhere d is the feature dimension. The scaling factor d1 prevents the dot products from growing too large, which would push the softmax into regions with vanishing gradients. In this simplified block, the queries are the latents, and both keys and values are formed by concatenating the image features X and the latents along the token (row) axis.
A residual connection is then applied: the attention output is added back to the original latents. This stabilizes training and ensures that information from the original latents is preserved. The "add-and-project" phrasing in the problem refers to this residual addition followed by the implicit projection that occurs through the attention mechanism itself (since no separate linear projections are used in this simplified version).
2. Algorithm Approach
The implementation follows a direct matrix-computation pipeline:
- Concatenate the image features X and the latents along the row axis to form the key/value matrix KV of shape (N+Q)×d.
- Compute attention scores: Multiply the latents (Q×d) by the transpose of KV (d×(N+Q)) and scale by d.
- Apply numerically stable softmax row-wise over the attention scores.
- Weighted sum: Multiply the softmax weights by KV to get the attention output of shape Q×d.
- Residual add: Add the attention output to the original latents.
- Round each element to 4 decimal places.
This is a single forward pass of a simplified transformer block with no separate query/key/value projection matrices.
3. Step-by-Step Strategy
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
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.