2D Sine-Cosine Positional Embedding
Problem Statement
Build the fixed 2D sin-cos positional embedding used by many ViT/MAE implementations: concatenate an independent 1D sin-cos embedding for the row coordinate and the column coordinate.
Background
For a 1D position pos and an embedding width d (even), the sin-cos embedding uses d/2 frequencies
ωi=100002i/d1,i=0,…,d/2−1
and emits [sin(posw_0), ..., sin(posw_{d/2-1}), cos(posw_0), ..., cos(posw_{d/2-1})].
For a 2D grid, each patch has a (row, col) coordinate. Give half the width to the row embedding and half to the column embedding, then concatenate [row_emb, col_emb] into a length-dim vector. Patches are enumerated row-major.
Your Task
Implement:
def pos_embed_2d(grid, dim):
Return a nested list of shape (grid*grid, dim), every value rounded to 4 decimals.
Input Format
- grid (int): side of the square patch grid.
- dim (int): total embedding width, divisible by 4.
Output Format
- A (grid*grid, dim) nested list, rounded to 4 decimals.
Sample
emb = pos_embed_2d(2, 4)
print(emb[0])
Output:
[0.0, 1.0, 0.0, 1.0]
Example:
emb = pos_embed_2d(2, 4) print(emb[0])
[0.0, 1.0, 0.0, 1.0]
- Split the total dimension dim=4 into two halves for the row and column embeddings, so each 1D embedding has width daxis=4/2=2.
- Determine the coordinates for the first patch (index 0) in the 2×2 grid: it is at row r=0 and column c=0.
- Calculate the row embedding for position 0 with width 2. The single frequency is ω0=10000−2(0)/2=1. The angles are 0⋅1=0, yielding [sin(0),cos(0)]=[0.0,1.0].
- Calculate the column embedding for position 0 with width 2. Similarly, the angle is 0, yielding [sin(0),cos(0)]=[0.0,1.0].
- Concatenate the row and column embeddings to form the final vector: [0.0,1.0]⊕[0.0,1.0]=[0.0,1.0,0.0,1.0].
- The final output is [0.0, 1.0, 0.0, 1.0]
Constraints:
dimis divisible by 4; each axis getsdim/2, each withdim/4frequencies.- Position 0 gives sin=0, cos=1, so the first patch is
[0,1,...,0,1]. - Round every entry to 4 decimals; avoid emitting
-0.0.
1. Background Knowledge
Positional embeddings are a critical component of Vision Transformers (ViT) because the self-attention mechanism is permutation-invariant; without explicit position information, the model cannot distinguish which patch is at the top-left versus the bottom-right of an image. While learned embeddings are common, sinusoidal (sin-cos) embeddings offer a fixed, parameter-free alternative that generalizes well to resolutions different from training.
The 1D sinusoidal encoding for a position pos with embedding width d (where d is even) uses d/2 distinct frequencies. The frequency for the i-th dimension is defined as:
ωi=100002i/d1The resulting vector is constructed by concatenating the sine terms followed by the cosine terms:
[sin(pos⋅ω0),…,sin(pos⋅ωd/2−1),cos(pos⋅ω0),…,cos(pos⋅ωd/2−1)]This structure ensures that the embedding is smooth and that nearby positions have similar embeddings, while distant positions are more distinguishable.
For a 2D grid, the standard approach (used in MAE and many ViT variants) is to treat the row and column coordinates independently. The total embedding dimension dim is split in half: the first dim/2 dimensions encode the row position, and the last dim/2 dimensions encode the column position. These two 1D embeddings are then concatenated to form the final 2D positional vector.
2. Algorithm Approach
The solution follows a decomposition and composition pattern:
- Decompose the 2D problem into two independent 1D problems (row and column).
- Generate 1D sinusoidal embeddings for each coordinate using the specified frequency formula.
- Compose the final embedding by concatenating the row and column vectors.
- Iterate over all patches in row-major order to build the full embedding matrix.
The core mathematical operation is vectorized frequency generation. For a given position p and half-width h=d/2, you compute the angles p⋅ωi for i=0,…,h−1, then apply sin and cos to these angles.
3. Step-by-Step Strategy
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
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.