Implement 2D Positional Embeddings

Difficulty: 🟡 Medium

Companies: Anthropic, DeepMind, Midjourney, Runway

Problem Statement

Vision Transformers (ViT) divide an image into a grid of patches and process them as a sequence. Unlike text, these patches have two-dimensional spatial structure (row and column). Standard 1D sinusoidal positional embeddings lose this 2D spatial information.

Your task is to implement 2D sinusoidal positional embeddings that encode both row and column position using sine/cosine functions. The embedding dimension is split in half: the first half encodes row position and the second half encodes column position, each using the standard sinusoidal formula.

Requirements

1. create_2d_sinusoidal_embeddings(height, width, d_model) — Returns a tensor of shape (height * width, d_model) containing the 2D positional embeddings.
2. Dimension split — First d_model // 2 dimensions encode row position; last d_model // 2 dimensions encode column position.
3. Sinusoidal formula — For each half, use alternating sine and cosine: PE(pos, 2i) = sin(pos / 10000^(2i/d_half)), PE(pos, 2i+1) = cos(pos / 10000^(2i/d_half)).

Constraints

• ✅ Must work for any height, width, and even d_model.
• ✅ Embeddings at neighboring grid positions should be more similar (higher cosine similarity) than distant positions.
• ✅ Row embeddings and column embeddings should be independent.
• ❌ Do not use nn.Embedding or learned embeddings.

• Compute the frequency divisors: div_term = torch.exp(torch.arange(0, d_half, 2) * -(math.log(10000.0) / d_half)).
• Create position vectors: row_pos = torch.arange(height).unsqueeze(1) and similarly for columns.
• Fill even indices with sine, odd indices with cosine for each half.
• Concatenate the row and column embeddings, then reshape to (height * width, d_model).

Implement

create_2d_sinusoidal_embeddings

Companies

Anthropic, DeepMind, Midjourney, Runway