NTK-Aware RoPE Base Scaling

Problem

NTK-aware RoPE scaling raises the base frequency to stretch context length without retraining: new_base = theta_base * scale^(dim / (dim - 2)), where scale is the context-extension factor and the exponent comes from spreading the interpolation over dimension pairs.

The signature is ntk_rope_base(theta_base, dim, scale). Return a float.

Hint

A scale factor of 1 leaves the base unchanged; larger context-extension factors raise it.

Test cases

case 1

in: [10000,64,2]

expected: 20452.2287120254

case 2

in: [10000,64,1]

expected: 10000

case 3

in: [10000,8,4]

expected: 63496.042078728

case 4

in: [500000,128,4]

expected: 2044497.121624311

Starter codePython

def ntk_rope_base(theta_base, dim, scale):
    # Your code here
    pass

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