Learned Position Embeddings › Extrapolation: Beyond Training Length

Michael BrenndoerferUpdated June 1, 20255 min read
Language AI HandbookAll chapters
Part 14 of 67Positional Encoding
Chapter 98Learned Position Embeddings › Extrapolation: Beyond Training Length
Reading section 8 of 19 · Extrapolation: Beyond Training Length · 5 minSection 8 of 19 · 5 min

Book contents

Section 8 of 19 · 5 min read

Learned Position Embeddings › Extrapolation: Beyond Training LengthOne section · Sections in this chapter8 / 19 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

What happens if you need to process sequences longer than Lmax⁡L_{\max}? With learned embeddings, you have several options, none of which are ideal.

Option 1: Extend and fine-tune. Add new position embeddings for positions beyond Lmax⁡L_{\max}, initialize them (perhaps by interpolating from existing embeddings), and fine-tune on longer sequences. This works but requires additional training. The existing embeddings for positions 0 through Lmax⁡−1L_{\max} - 1 remain intact, but the new embeddings need training data that includes long sequences to learn meaningful representations for the extended positions. This approach is used in practice when a model trained on short sequences needs to handle longer documents: extend the table, fine-tune on long-document data, and the model adapts.

Option 2: Position interpolation. If you trained with Lmax⁡=512L_{\max} = 512 but need to handle 1024 tokens, you can interpolate the position indices. Position 512 in the long sequence uses the embedding for position 256 in the original table. This trick, called position interpolation, works surprisingly well for moderate extensions.

In[17]:
Code
def position_interpolation(pos_embed, target_len, original_max_len):
    """
    Interpolate position embeddings for longer sequences.

    Maps positions [0, target_len) to [0, original_max_len) and uses
    linear interpolation between the two nearest positions.
    """
    # Scale factor: how much to compress positions
    scale = original_max_len / target_len

    interpolated = np.zeros((target_len, pos_embed.embed_dim))

    for i in range(target_len):
        # Map to position in original table
        orig_pos = i * scale

        # Linear interpolation between floor and ceil positions
        low_pos = int(np.floor(orig_pos))
        high_pos = min(int(np.ceil(orig_pos)), original_max_len - 1)

        if low_pos == high_pos:
            interpolated[i] = pos_embed.embeddings[low_pos]
        else:
            # Weight based on fractional position
            weight = orig_pos - low_pos
            interpolated[i] = (1 - weight) * pos_embed.embeddings[
                low_pos
            ] + weight * pos_embed.embeddings[high_pos]

    return interpolated


# Extend from 128 to 256 positions using interpolation
original_max = 128
target_len = 256
interpolated_embeds = position_interpolation(
    pos_embed, target_len, original_max
)
Out[18]:
Console
Position Interpolation:
  Original max length: 128
  Target length:       256
  Scale factor:        0.50

Example mappings:
  New position   0 -> Original position 0.0
  New position  64 -> Original position 32.0
  New position 128 -> Original position 64.0
  New position 192 -> Original position 96.0
  New position 255 -> Original position 127.5

Position interpolation effectively "stretches" the position embedding table. The model sees positions at half the resolution, but the relative ordering is preserved. This works because the model primarily cares about relative positions, and those relationships are maintained under scaling. Research on extending GPT-2's context from 1024 to 2048 using interpolation found surprisingly small degradation in language modeling performance, suggesting that the absolute scale of position representations matters less than their relative structure.

The mathematical intuition is straightforward. If position embeddings encode a smooth manifold in the embedding space (which they do after training), then interpolated points along that manifold are reasonable approximations of what trained embeddings at those positions would look like. The interpolation maintains the smooth variation property that makes position embeddings effective.

Let's visualize how position interpolation affects the embedding structure:

In[19]:
Code
# Compare similarity structure before and after interpolation
# First, create a simulated trained embedding for the original table
original_trained = create_trained_position_embeddings(128, 64, seed=42)


# Interpolate to 256 positions
class SimpleEmbed:
    def __init__(self, embeddings):
        self.embeddings = embeddings
        self.embed_dim = embeddings.shape[1]


interpolated_256 = position_interpolation(
    SimpleEmbed(original_trained), 256, 128
)

# Compute similarity matrices for both
sim_original = cosine_similarity_matrix(original_trained)
sim_interpolated = cosine_similarity_matrix(interpolated_256)
Out[20]:
Visualization
Heatmap of 64x64 position similarities with bright diagonal band.
Original position similarity matrix for 128 positions. The diagonal band of high similarity, gradually fading with distance, shows the smooth positional geometry learned during training.
Heatmap of 128x128 interpolated position similarities with similar but stretched pattern.
Interpolated position similarity matrix for 256 positions. The same structural pattern is preserved but stretched across twice as many positions. Each original position maps to two new positions, and the relative relationships between all pairs are maintained.

The side-by-side comparison shows that position interpolation preserves the essential structure: nearby positions remain similar, and similarity decays with distance. The interpolated version has twice as many positions, but the relative relationships are maintained. This explains why interpolation works reasonably well for moderate length extensions.

Option 3: Sliding window. Process long sequences in overlapping chunks of length Lmax⁡L_{\max}. Each chunk gets proper position embeddings within the window (positions 0 through Lmax⁡−1L_{\max} - 1), but the model cannot attend across window boundaries. Predictions near the boundary of each chunk may be less accurate because they lack full context. This is common for very long documents where sliding window processing is acceptable.

The extrapolation problem is a significant limitation of learned position embeddings. Models trained with shorter contexts may struggle when forced to process longer sequences, even with interpolation tricks. This limitation has motivated a body of research into position encodings that generalize better to unseen lengths, including relative position encodings (which encode the distance between positions rather than their absolute locations) and rotary position embeddings (RoPE, which encodes position by rotating query and key vectors). We'll explore these approaches in later chapters.