Positional Embedding Interpolation

Problem

Linearly interpolate a positional embedding table from old_len rows to new_len rows. For output position i, the source position is i * (old_len - 1) / (new_len - 1); blend the two neighboring rows with the fractional part. If old_len is 1 the table repeats; if new_len is 1 the first row is returned.

The signature is interpolate_positional_embedding(pe, new_len).

Hint

Interpolation lets a model trained at one sequence length run at another.

Test cases

case 1

in: [[[0,0],[1,1],[2,2]],5]

expected: [[0,0],[0.5,0.5],[1,1],[1.5,1.5],[2,2]]

case 2

in: [[[0],[10]],3]

expected: [[0],[5],[10]]

case 3

in: [[[0,0],[1,2],[2,4],[3,6]],2]

expected: [[0,0],[3,6]]

case 4

in: [[[5,7]],3]

expected: [[5,7],[5,7],[5,7]]

Starter codePython

import math
def interpolate_positional_embedding(pe, new_len):
    # Your code here
    pass

4 test cases · run your code against them in the editor