Pixel-Shuffle Token Downsampling

mediumPROPatch Embedding & Visual TokensVision-Language ModelsVLM 1: Vision-Language Models

Problem Statement

Many recent VLMs (InternVL, some LLaVA variants) shrink the number of visual tokens with a pixel-shuffle (space-to-depth): a grid x grid feature map of width C becomes a (grid/r) x (grid/r) map of width Crr by folding each r x r spatial block into the channel dimension. Report the resulting token count and channel width, and the compression ratio.

Background

Attention cost is quadratic in token count, so halving the grid side (r=2) cuts tokens 4x while preserving information by widening channels 4x. The new token count is (grid // r) ** 2 and the new channel width is C * r * r. The token compression ratio is r*r.

Your Task

Implement:

def pixel_shuffle_tokens(grid, channels, r):

Return a dict with "tokens", "channels", "ratio" (all ints). grid is divisible by r.

Input Format

grid (int), channels (int), r (int); grid % r == 0.

Output Format

A dict of three ints.

Sample

print(pixel_shuffle_tokens(24, 1024, 2))

Output:

{'tokens': 144, 'channels': 4096, 'ratio': 4}

Example

Input

print(pixel_shuffle_tokens(24, 1024, 2))

Output

{'tokens': 144, 'channels': 4096, 'ratio': 4}

Reasoning

The spatial grid is reduced by the shuffle factor $r$ to determine the new side length: $24 / 2 = 12$ .
The new token count is calculated by squaring the reduced side length, representing the total number of visual tokens: $12^2 = 144$ .
The channel width is expanded by the square of the shuffle factor to preserve the total information volume: $1024 \times 2^2 = 4096$ .
The compression ratio is the square of the shuffle factor, indicating how many original tokens are merged into one: $2^2 = 4$ .
The final output is {'tokens': 144, 'channels': 4096, 'ratio': 4}

Constraints

grid % r == 0, 1 <= r <= grid.
tokens = (grid // r) ** 2, channels = channels_in * r * r, ratio = r * r.
Return ints.

Learn this in

Benchmarks, Deployment & the FrontierVLM 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

In Vision-Language Models (VLMs), images are converted into sequences of visual tokens for processing by a transformer. A common method is patch embedding, where an image is divided into non-overlapping patches, and each patch becomes a single token. If an image is resized to a $grid \times grid$ resolution, the initial number of visual tokens is $grid^2$ . Because the self-attention mechanism in transformers has a computational complexity of $O(N^2)$ with respect to the sequence length $N$ , reducing the number of tokens is critical for efficiency.

Pixel-shuffle (also known as space-to-depth) is a technique used to downsample spatial dimensions while preserving information by expanding the channel dimension. Instead of simply discarding pixels, it rearranges the data. For a feature map of shape $(H, W, C)$ , applying a pixel-shuffle with a factor $r$ groups every $r \times r$ block of spatial pixels into the channel dimension. The spatial dimensions are divided by $r$ , and the channel dimension is multiplied by $r^2$ . This operation is bijective; no information is lost, but the representation changes from many low-dimensional tokens to fewer high-dimensional tokens.

This trade-off is particularly useful in VLMs like InternVL or LLaVA variants. By increasing $r$ , the number of tokens sent to the language model decreases quadratically, significantly reducing attention costs. For example, if $r=2$ , the token count is reduced by a factor of 4, but each token now carries 4 times the channel information. This allows the model to maintain representational capacity while improving computational throughput.

2. Algorithm Approach

The problem requires implementing a mathematical transformation rather than manipulating actual tensor data. The approach involves:

Understanding the Transformation: Recognize that pixel-shuffle with factor $r$ transforms a grid of size $G$ into a new grid size of $G/r$ .
Calculating New Dimensions:

New spatial dimension: $G_{new} = G / r$ .
New token count: $N_{new} = G_{new}^2 = (G/r)^2$ .
New channel width: $C_{new} = C \times r^2$ .

Computing Compression Ratio: The ratio is the factor by which the token count is reduced, which is $r^2$ .
Returning Results: Package these integer values into a dictionary.

Since the problem guarantees that grid is divisible by r, integer division is safe and exact.

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