YaRN › Visualizing YaRN's Effect

def compute_attention_matrix(
    seq_len, d_model, freqs, scale_attention=1.0, seed=42
):
    """Compute full attention weight matrix."""
    rng = np.random.default_rng(seed)

    Q = rng.standard_normal((seq_len, d_model))
    K = rng.standard_normal((seq_len, d_model))

    Q_rotated = np.array([apply_rope(Q[i], i, freqs) for i in range(seq_len)])
    K_rotated = np.array([apply_rope(K[i], i, freqs) for i in range(seq_len)])

    scores = Q_rotated @ K_rotated.T / np.sqrt(d_model)
    scores = scores * scale_attention

    weights = softmax(scores)
    return weights


# Compute attention patterns for each method
seq_len = 48
d_model = 64
scale = 16.0

# Recompute extension freqs at this scale so differences are pronounced
_pi_freqs_vis = compute_pi_frequencies(d_model, scale=scale)
_ntk_freqs_vis = compute_ntk_frequencies(d_model, scale=scale)

# YaRN with temperature
yarn_module = YaRNRoPE(d_model=d_model, original_context=2048, scale=scale)
yarn_attn = compute_attention_matrix(
    seq_len,
    d_model,
    yarn_module.adjusted_freqs,
    scale_attention=np.sqrt(yarn_module.temperature),
)

# Original (no extension)
original_attn = compute_attention_matrix(seq_len, d_model, original_freqs)

# Position Interpolation
pi_attn = compute_attention_matrix(seq_len, d_model, _pi_freqs_vis)

# NTK-aware
ntk_attn = compute_attention_matrix(seq_len, d_model, _ntk_freqs_vis)