Attention Pooling to a Single Token

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

Problem Statement

The simplest bridge is attention pooling: collapse many patch features into one vector using a set of attention weights. Compute the weighted average of the patch features.

Background

Given patch features X (N x D) and attention weights w (length N, already summing to 1), the pooled token is

$p = \sum_{i} w_i\, X_i$

a convex combination of the rows. This is exactly what a one-query attention layer produces after its softmax.

Your Task

Implement:

def attention_pool(X, w):

Return the pooled D-vector as a list rounded to 4 decimals.

Input Format

X: N x D nested list.
w: list of N weights summing to 1.

Output Format

A list of D floats rounded to 4 decimals.

Sample

print(attention_pool([[1.0, 2.0], [3.0, 4.0]], [0.5, 0.5]))

Output:

[2.0, 3.0]

Example

Input

print(attention_pool([[1.0, 2.0], [3.0, 4.0]], [0.5, 0.5]))

Output

[2.0, 3.0]

Reasoning

Identify the patch features $X$ as a $2 \times 2$ matrix with rows $[1.0, 2.0]$ and $[3.0, 4.0]$ , and the attention weights $w$ as $[0.5, 0.5]$ .
Compute the pooled value for the first dimension by taking the weighted sum of the first column: $0.5 \times 1.0 + 0.5 \times 3.0 = 0.5 + 1.5 = 2.0$ .
Compute the pooled value for the second dimension by taking the weighted sum of the second column: $0.5 \times 2.0 + 0.5 \times 4.0 = 1.0 + 2.0 = 3.0$ .
Combine these results into the vector $[2.0, 3.0]$ and round each element to 4 decimal places, which leaves the values unchanged.
The final output is [2.0, 3.0]

Constraints

1 <= N <= 5000, 1 <= D <= 4096.
w is a valid distribution over the N rows.
Round every entry to 4 decimals; avoid -0.0.

Learn this in

Introduction to Vision-Language ModelsVLM Study Plan Vision-Language AlignmentVLM 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

Attention pooling is the simplest form of a bridge module used in Vision-Language Models (VLMs) to reduce a sequence of patch features into a single representation. In a standard Vision Transformer, an image is split into $N$ patches, each encoded into a $D$ -dimensional vector. To feed this into a language model that expects a single token, we compute a weighted average of all patch features. This operation is mathematically equivalent to a one-query attention mechanism: the query attends to all keys, and the softmax-normalized attention scores become the weights $w_i$ .

The core operation is a convex combination. Given patch features $X \in \mathbb{R}^{N \times D}$ and weights $w \in \mathbb{R}^N$ where $\sum_{i=0}^{N-1} w_i = 1$ , the pooled vector is:

p = \sum_{i=0}^{N-1} w_i X_i

Each element $p_j$ of the output is the weighted sum of the $j$ -th column of $X$ . This is a linear operation that preserves the dimensionality $D$ while collapsing the sequence length $N$ to 1.

In the context of Q-Former and other bridge modules, attention pooling serves as the baseline. More complex bridges use cross-attention with learnable queries, but the fundamental idea remains: project multiple visual tokens into a compact representation that the language model can consume.

2. Algorithm Approach

This is a straightforward weighted sum computation. The approach is:

Initialize an accumulator vector of length $D$ with zeros.
For each patch $i$ (from $0$ to $N-1$ ), multiply the entire row $X[i]$ by the scalar weight $w[i]$ .
Add the scaled row to the accumulator.
After processing all patches, round each element to 4 decimal places.

This can be implemented with nested loops or vectorized operations. Since the problem specifies Python lists, a nested loop approach is natural and efficient enough for typical input sizes.

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