File size: 6,464 Bytes

d431cc8

## 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})`:

$$
\begin{pmatrix} x'_{2i} \\ x'_{2i+1} \end{pmatrix} =
\begin{pmatrix} \cos(m\theta_i) & -\sin(m\theta_i) \\ \sin(m\theta_i) & \cos(m\theta_i) \end{pmatrix}
\begin{pmatrix} x_{2i} \\ x_{2i+1} \end{pmatrix}
$$

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$$
$$K = W_K \cdot x + b_K$$

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

```python
# 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} = q_{2i} \cdot \cos(m\theta_i) - q_{2i+1} \cdot \sin(m\theta_i)
$$

Expanding with bias:
$$
q_{2i} = (W_Q \cdot x)_{2i} + b_{2i}
$$

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

```python
# 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} = \underbrace{(W_Q x)_{2i}}_{O(1)} \cdot \cos(m\theta_i) + \underbrace{b_{2i}}_{O(100)} \cdot \cos(m\theta_i) - \underbrace{(W_Q x)_{2i+1}}_{O(1)} \cdot \sin(m\theta_i) - \underbrace{b_{2i+1}}_{O(100)} \cdot \sin(m\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: \quad q'_{2i} \approx -q_{2i+1} \cdot \sin(m\theta_i) - b_{2i+1} \cdot \sin(m\theta_i)
$$

The relative error becomes:
$$
\epsilon_{rel} = \frac{|q'_{TT} - q'_{HF}|}{|q'_{HF}|} \propto \frac{\epsilon_{bf16}}{|\sin(m\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{QK^T}{\sqrt{d_k}}\right) V
$$

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

$$
QK^T = Q \cdot [K_0 + \epsilon_0, K_1 + \epsilon_1, ..., K_{127} + \epsilon_{127}]^T
$$

The errors accumulate in the dot product:
$$
\sum_{j=0}^{127} \epsilon_j \cdot Q \approx O(\sqrt{128}) \cdot \epsilon_{avg} \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} = (W_Q x)_{2i} \cdot \cos(m\theta_i) - (W_Q x)_{2i+1} \cdot \sin(m\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{|b_{max}|^2}{|W_Q x|^2} \cdot \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.