Perceiver Resampler Forward Pass

Problem Statement

Implement one multi-head Perceiver Resampler block: a fixed set of learned latent tokens cross-attends to image features, with the latents themselves concatenated onto the keys and values, followed by an output projection and a residual connection.

Background

Flamingo's Perceiver Resampler turns a variable-length grid of visual features into a fixed number of tokens. It is not plain cross-attention - the paper's block builds keys and values from the image features with the latents appended:

$K = [X_f;\, L]\,W_K, \qquad V = [X_f;\, L]\,W_V, \qquad Q = L\,W_Q$

so the latents can attend to each other as well as to the image, mixing information across the latent set within a single block. The key/value length is therefore n_feats + n_latents, while the query length - and so the output length - stays n_latents.

The rest is standard multi-head attention. Split the d-wide projections into h heads of width d_h = d / h, run scaled dot-product attention per head (scaling by 1/sqrt(d_h), the head width, not the model width), concatenate the heads back to d, apply W_O, and add the residual:

$\text{out} = \text{Concat}(\text{head}_1..\text{head}_h)\,W_O + L$

Head splitting is a reshape of the last axis into (h, d_h): head i owns columns **id_h ... (i+1)d_h - 1, contiguously.

Your Task

Implement:

def resampler_forward(latents, feats, Wq, Wk, Wv, Wo, num_heads):

Return the (n_latents, d) output as a nested list, every entry rounded to 4 decimals.

Input Format

• latents - nested list, shape (n_latents, d)
• feats - nested list, shape (n_feats, d)
• Wq, Wk, Wv, Wo - nested lists, each shape (d, d), applied on the right (X @ W)
• num_heads - integer dividing d

Output Format

A nested list of n_latents rows of d floats, each rounded to 4 decimals.

Sample

L = [[1.0, 0.0], [0.0, 1.0]]
X = [[1.0, 1.0], [2.0, 0.0]]
I2 = [[1.0, 0.0], [0.0, 1.0]]
print(resampler_forward(L, X, I2, I2, I2, I2, 1))

Output:

[[2.3395, 0.3302], [0.8302, 1.6698]]