One Perceiver Resampler Block

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

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

Input

latents = [[1.0, 0.0]]
X = [[0.0, 2.0]]
print(perceiver_block(latents, X))

Output

[[1.6698, 0.6605]]

Reasoning

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.

Learn this in

Multimodal LLMsLLM 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

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:

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V

where $d$ is the feature dimension. The scaling factor $\frac{1}{\sqrt{d}}$ 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) \times d$ .
Compute attention scores: Multiply the latents ( $Q \times d$ ) by the transpose of $KV$ ( $d \times (N+Q)$ ) and scale by $\sqrt{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 \times 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

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