YaRN context-window scaling: NTK-by-parts + ramp + temperature

Peng, Quesnelle, Fan and Shippole (2023), "YaRN: Efficient Context Window Extension of Large Language Models", arXiv:2309.00071 (ICLR 2024). Eq 17 rotation counts r(d) = L theta_d / (2 pi); Eq 18 ramp gamma; Eq 20 blend h(theta_d) = (1 - gamma) theta_d / s + gamma theta_d; Eq 22 temperature sqrt(1/t) = 0.1 ln(s) + 1.

Usage

Yarn(base, s, d, L, beta_fast = 32, beta_slow = 1)

Arguments

base

RoPE base b (scalar > 1), or the d/2 frequencies directly (vector).

s

Context extension scale factor (> 0).

d

Head embedding width, positive even integer.

L

Original context length (> 0).

beta_fast, beta_slow

Rotation-count bounds (paper Eq 18 with alpha = beta_slow, beta = beta_fast; defaults 32 and 1).

Value

List with theta, theta_new, rotations, gamma, temperature, logit_scale, scale, estimate, n, method.

References

Peng, B., Quesnelle, J., Fan, H. and Shippole, E. (2023), arXiv:2309.00071, Sections 3.2-3.4, Eqs 17/18/20/22.

Examples

y <- Yarn(base = 10000, s = 4, d = 64, L = 4096)
y$temperature
#> [1] 0.7713211