code/models/common/tests/modules/attention/low_pcc_notes.md

Mathematical Explanation of the PCC Degradation

1. The RoPE Transformation

Rotary Position Embedding (RoPE) applies a rotation to Q and K vectors based on position. For a query/key vector at position m, RoPE multiplies pairs of elements by complex rotation:

For head dimension d, split into pairs (x_{2i}, x_{2i+1}):

where $\theta_i = 10000^{-2i/d}$

2. The HuggingFace vs Meta Format

HuggingFace format stores Q/K as:

[r₀, r₁, r₂, ..., r₆₃, i₀, i₁, i₂, ..., i₆₃]  (first half = "real", second half = "imag")

Meta/TTNN format (after _reverse_permute) stores as:

[r₀, i₀, r₁, i₁, r₂, i₂, ..., r₆₃, i₆₃]  (interleaved pairs)

RoPE operates on adjacent pairs, so Meta format aligns directly with how RoPE computes rotations.

3. Where Q/K Bias Causes Issues

For Qwen2.5-7B, the Q/K projections have biases:

$$Q=W_Q\cdot x+b_{Q}Q\; =\; W\_Q\; \backslash cdot\; x\; +\; b\_Q$$ $$K=W_K\cdot x+b_{K}K\; =\; W\_K\; \backslash cdot\; x\; +\; b\_K$$

The bias b_Q is stored in HuggingFace format. When we apply _reverse_permute_1d:

# HF: b = [b_r0, b_r1, ..., b_r63, b_i0, b_i1, ..., b_i63]
# After _reverse_permute_1d:
# Meta: b' = [b_r0, b_i0, b_r1, b_i1, ..., b_r63, b_i63]

4. The Numerical Precision Problem

The RoPE rotation for element pair $(q_{2i}, q_{2i+1})$ at position $m$ is:

$$q_{2i}^{\mathrm{\prime }}=q_{2i}\cdot \cos (m{\theta }_i)-q_{2i+1}\cdot \sin (m{\theta }_i)q\text{'}\_\{2i\}\; =\; q\_\{2i\}\; \backslash cdot\; \backslash cos(m\backslash theta\_i)\; -\; q\_\{2i+1\}\; \backslash cdot\; \backslash sin(m\backslash theta\_i)$$

Expanding with bias: $$q_{2i}=(W_Q\cdot x{)}_{2i}+b_{2i}q\_\{2i\}\; =\; (W\_Q\; \backslash cdot\; x)\_\{2i\}\; +\; b\_\{2i\}$$

The issue is that the bias values for Qwen2.5-7B are extremely large:

# Qwen2.5-7B bias statistics (layer 0):
# Q Bias: min=-48.25, max=46.25, std=2.97, abs_max=48.25
# K Bias: min=-164.0, max=171.0, std=26.75, abs_max=171.0  # Very large!
# V Bias: min=-1.53, max=2.58, std=0.19, abs_max=2.58

The K bias can be as large as 171.0 - this is enormous compared to typical activation magnitudes of O(1).

When RoPE rotates, the computation becomes:

$$q_{2i}^{\mathrm{\prime }}=\underset{O(1)}{\underbrace{(W_{Q}x{)}_{2i}}}\cdot \cos (m{\theta }_i)+\underset{O(100)}{\underbrace{b_{2i}}}\cdot \cos (m{\theta }_i)-\underset{O(1)}{\underbrace{(W_{Q}x{)}_{2i+1}}}\cdot \sin (m{\theta }_i)-\underset{O(100)}{\underbrace{b_{2i+1}}}\cdot \sin (m{\theta }_i)q\text{'}\_\{2i\}\; =\; \backslash underbrace\{(W\_Q\; x)\_\{2i\}\}\_\{O(1)\}\; \backslash cdot\; \backslash cos(m\backslash theta\_i)\; +\; \backslash underbrace\{b\_\{2i\}\}\_\{O(100)\}\; \backslash cdot\; \backslash cos(m\backslash theta\_i)\; -\; \backslash underbrace\{(W\_Q\; x)\_\{2i+1\}\}\_\{O(1)\}\; \backslash cdot\; \backslash sin(m\backslash theta\_i)\; -\; \backslash underbrace\{b\_\{2i+1\}\}\_\{O(100)\}\; \backslash cdot\; \backslash sin(m\backslash theta\_i)$$

The bias terms dominate the computation, making numerical precision errors much more significant.

5. Position-Dependent Amplification

The key insight is that $\cos(m\theta_i)$ and $\sin(m\theta_i)$ vary dramatically across positions:

• For small $i$ (low frequencies): $\theta_i \approx 1$, so $\cos(m\theta_i)$ oscillates rapidly
• For large $i$ (high frequencies): $\theta_i \approx 10^{-4}$, so $\cos(m\theta_i) \approx 1$

At certain positions, the combination of:

1. Large bias values
2. Near-zero $\cos$ or $\sin$ (causing subtraction of nearly equal large numbers)
3. bfloat16 precision (only 7 bits mantissa)

Causes catastrophic cancellation:

$$\text{When }\cos (m{\theta }_i)\approx 0:q_{2i}^{\mathrm{\prime }}\approx -q_{2i+1}\cdot \sin (m{\theta }_i)-b_{2i+1}\cdot \sin (m{\theta }_i)\backslash text\{When\; \}\; \backslash cos(m\backslash theta\_i)\; \backslash approx\; 0:\; \backslash quad\; q\text{'}\_\{2i\}\; \backslash approx\; -q\_\{2i+1\}\; \backslash cdot\; \backslash sin(m\backslash theta\_i)\; -\; b\_\{2i+1\}\; \backslash cdot\; \backslash sin(m\backslash theta\_i)$$

The relative error becomes: $${ϵ}_{rel}=\frac{\mathrm{\mid }{q}_{TT}^{\mathrm{\prime }}-q_{HF}^{\mathrm{\prime }}\mathrm{\mid }}{\mathrm{\mid }{q}_{HF}^{\mathrm{\prime }}\mathrm{\mid }}\propto \frac{{ϵ}_{bf16}}{\mathrm{\mid }\sin (m{\theta }_i)\mathrm{\mid }}\backslash epsilon\_\{rel\}\; =\; \backslash frac\{|q\text{'}\_\{TT\}\; -\; q\text{'}\_\{HF\}|\}\{|q\text{'}\_\{HF\}|\}\; \backslash propto\; \backslash frac\{\backslash epsilon\_\{bf16\}\}\{|\backslash sin(m\backslash theta\_i)|\}$$

When $\sin(m\theta_i)$ is small, the relative error blows up.

6. Why Position 7 in TTTv1 and Varying Positions in TTTv2?

Looking at $\theta_i = 10000^{-2i/128}$ for $i=0$:

• $\theta_0 = 1$
• At position $m=7$: $\cos(7) \approx 0.754$, $\sin(7) \approx 0.657$

For $i=1$: $\theta_1 = 10000^{-1/64} \approx 0.891$

• At position $m=7$: $\cos(7 \cdot 0.891) \approx \cos(6.24) \approx 0.998$

The specific positions where PCC drops depend on which frequency components have near-zero $\cos/\sin$ values, combined with which bias elements have the largest magnitudes.

7. Why Prefill+Decode is Worse (0.956 vs 0.972)

In SDPA (Scaled Dot-Product Attention):

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V\backslash text\{Attention\}(Q,\; K,\; V)\; =\; \backslash text\{softmax\}\backslash left(\backslash frac\{QK^T\}\{\backslash sqrt\{d\_k\}\}\backslash right)\; V$$

With 128 prefilled K entries, each with small numerical error $\epsilon_k$:

$$Q{K}^T=Q\cdot [K_0+{ϵ}_0,K_1+{ϵ}_1,\mathrm{.}\mathrm{.}\mathrm{.},K_{127}+{ϵ}_{127}{]}^{T}QK^T\; =\; Q\; \backslash cdot\; [K\_0\; +\; \backslash epsilon\_0,\; K\_1\; +\; \backslash epsilon\_1,\; ...,\; K\_\{127\}\; +\; \backslash epsilon\_\{127\}]^T$$

The errors accumulate in the dot product: $$\sum _{j=0}^{127}{ϵ}_j\cdot Q\approx O(\sqrt{128})\cdot {ϵ}_{avg}\cdot \mathrm{\mid }\mathrm{\mid }Q\mathrm{\mid }\mathrm{\mid }\backslash sum\_\{j=0\}^\{127\}\; \backslash epsilon\_j\; \backslash cdot\; Q\; \backslash approx\; O(\backslash sqrt\{128\})\; \backslash cdot\; \backslash epsilon\_\{avg\}\; \backslash cdot\; ||Q||$$

This ~11x amplification of errors (from $\sqrt{128}$) explains why prefill+decode PCC (0.956) is lower than decode-only PCC (0.972).

8. Comparison with Other Models

To validate this analysis, we compare bias magnitudes across different models:

Model	Has Q/K Bias	K Bias abs_max	Q Bias abs_max	Test PCC
Qwen2.5-7B	Yes	171.0	48.25	0.97 (decode-only)
DeepSeek-R1-14B	Yes	21.75	12.0	0.98+
Llama-3.1-8B	No	N/A	N/A	0.99+

Why Llama-3.1-8B Has No Issues

Llama models have attention_bias=False - there are no Q/K biases at all. The RoPE rotation only operates on:

$$q_{2i}^{\mathrm{\prime }}=(W_{Q}x{)}_{2i}\cdot \cos (m{\theta }_i)-(W_{Q}x{)}_{2i+1}\cdot \sin (m{\theta }_i)q\text{'}\_\{2i\}\; =\; (W\_Q\; x)\_\{2i\}\; \backslash cdot\; \backslash cos(m\backslash theta\_i)\; -\; (W\_Q\; x)\_\{2i+1\}\; \backslash cdot\; \backslash sin(m\backslash theta\_i)$$

With typical activation magnitudes of O(1), the bfloat16 precision is sufficient and no catastrophic cancellation occurs.

Why DeepSeek-R1-14B Has Better PCC Than Qwen2.5-7B

DeepSeek-R1-14B does have Q/K biases, but they are ~8x smaller:

• K bias max: 21.75 (vs Qwen2.5-7B's 171.0)
• Q bias max: 12.0 (vs Qwen2.5-7B's 48.25)

Smaller biases mean:

1. The bias terms don't dominate the computation as much
2. Less catastrophic cancellation when $\cos/\sin$ approach zero
3. The relative error stays within acceptable bounds

The relationship between bias magnitude and PCC degradation is roughly:

$$\text{PCC degradation}\propto \frac{\mathrm{\mid }{b}_{max}{\mathrm{\mid }}^2}{\mathrm{\mid }{W}_{Q}x{\mathrm{\mid }}^2}\cdot {ϵ}_{bf16}\backslash text\{PCC\; degradation\}\; \backslash propto\; \backslash frac\{|b\_\{max\}|^2\}\{|W\_Q\; x|^2\}\; \backslash cdot\; \backslash epsilon\_\{bf16\}$$

With Qwen2.5-7B's K bias being ~8x larger, the PCC degradation is ~64x worse, explaining the observed difference.

Summary

The lower PCC for Qwen2.5-7B is caused by:

1. Large Q/K biases that get rotated by RoPE
2. Position-dependent $\cos/\sin$ values that can approach zero
3. bfloat16 precision limitations causing catastrophic cancellation when subtracting nearly-equal values
4. Error accumulation in SDPA when attending over many KV entries

This is a fundamental numerical precision characteristic of the model architecture, not a bug in TTTv1 or TTTv2.