Skip to content

add Q=K ablation results #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
teosssss wants to merge 2 commits into
base: main
Choose a base branch
from
Open

add Q=K ablation results #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5eb4c4b
add Q=K ablation results
Apr 9, 2026
11f40da
add cleaned WHT vs PCA three-way experiment
Apr 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 79 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,85 @@ By identifying these dimensions through a one-time **15-second calibration** and

4. **KV spectral asymmetry.** Keys: d_eff ≈ 4. Values: d_eff ≈ 40–55 (10–15× larger). This explains why low-rank compression fails for values while SQ succeeds.

### Representation Note

Throughout these experiments, "`1-bit keys`" should be read narrowly: one bit per rotated key dimension in the calibrated spectral basis, plus a separately stored vector norm. It is not a raw sign-only key in the original basis.

### Small-Model Sanity Checks

To check how much of the reported `1-bit keys` quality depends on the benchmark itself, I ran a reduced local MPS comparison on three small open models using the same compression configs but two different evaluation modes.

The two modes were:

1. `Proxy Q=K`
- queries are replaced with keys (`Q = K`)
- attention is evaluated with a non-causal self-attention proxy
- this is the more favorable setup

2. `Real Q + causal`
- queries are captured from the actual attention module during a forward pass
- the same compressed keys/values are evaluated with causal attention
- this is closer to the intended decoding setting

Important scope note:
- this is not a pure `Q=K` ablation, because the comparison changes two variables at once:
- query source: proxy `Q=K` vs true captured `Q`
- attention pattern: non-causal vs causal
- the result is still useful because it measures the gap between an optimistic proxy and a more realistic decoding-style evaluation
- but it does not by itself identify how much of the gap comes from `Q=K` specifically versus the causal mask

What stayed fixed between the two runs:
- same model
- same reduced coverage
- same sampled layers and heads
- same compression config
- same key/value reconstruction path

Coverage used for all rows below:
- `n_calib=8`
- `n_eval=4`
- `seq_len=512`
- `max_cached_vectors=256`
- sampled layers / sampled KV heads

Interpretation:
- `Proxy Q=K` and `Real Q + causal` have the same key/value reconstruction numbers for a given config.
- The gap is therefore caused by the evaluation setup, not by a different quantizer.
- A negative delta means the favorable proxy made the config look better than it does under real queries.

For the representative `1bit-K_6bit-V` setting, replacing the proxy with true captured queries and causal attention lowers attention cosine substantially:

| Model | Config | Proxy `Q=K` | Real Q + causal | Delta |
|---|---|---:|---:|---:|
| Qwen 2.5 1.5B | `1bit-K_6bit-V` | 0.7451 | 0.5250 | -0.2201 |
| TinyLlama 1.1B | `1bit-K_6bit-V` | 0.8644 | 0.4954 | -0.3690 |
| Qwen 2.5 0.5B | `1bit-K_6bit-V` | 0.8676 | 0.6070 | -0.2606 |

Files:
- `results/push_095/push095_Qwen-1.5B-proxy-small.json`
- `results/push_095/push095_Qwen-1.5B-realq-small.json`
- `results/push_095/push095_TinyLlama-1.1B-proxy-small.json`
- `results/push_095/push095_TinyLlama-1.1B-realq-small.json`
- `results/push_095/push095_Qwen-0.5B-proxy-small.json`
- `results/push_095/push095_Qwen-0.5B-realq-small.json`
- `results/push_095/qwen1p5b_1bit_claim_check_small.md`

Recommended follow-up for a cleaner attribution:
- run the full `2 x 2` evaluation matrix:
- `Q=K + non-causal`
- `Q=K + causal`
- `Real Q + non-causal`
- `Real Q + causal`
- that would isolate whether the optimism is driven more by proxy queries, by removing the causal mask, or by both
- the `Real Q + non-causal` setting has now been run for `Qwen 2.5 1.5B`
- file: `results/push_095/push095_Qwen-1.5B-realq-noncausal-small.json`
- for `1bit-K_6bit-V`, the three observed values are:
- `Q=K + non-causal`: `0.7451`
- `Real Q + non-causal`: `0.4124`
- `Real Q + causal`: `0.5250`
- on this reduced run, most of the optimism comes from the `Q=K` proxy rather than from removing the causal mask
- this means the current evidence supports the narrower statement "proxy queries are the main confound here"; the `Q=K + causal` corner is still missing if a full `2 x 2` attribution is needed

---

## Quick Start
Expand Down
216 changes: 216 additions & 0 deletions codex_wht_threeway_final_report.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,216 @@
# WHT vs PCA vs PolarQuant

This note summarizes the final cleaned-up version of Experiment 4 after fixing the calibration inconsistencies in the earlier drafts.

The final comparison is a three-way benchmark:

- `PCA`
- `WHT-only`
- `WHT PolarQuant-style`

The important cleanup is that `WHT-only` is now calibrated from **normalized** vectors, which makes it comparable to the PolarQuant branch.

## Goal

The point of this experiment is to separate three different questions that were previously mixed together:

1. What happens if we replace the learned PCA basis with a fixed WHT basis?
2. What happens if we keep WHT but use per-coordinate empirical calibration?
3. What happens if we keep WHT and instead use PolarQuant-style shared theoretical codebooks?

This gives a cleaner decomposition than the earlier two-way experiments.

## Variants

### PCA

- learn per-head K/V covariance from calibration data
- eigendecompose covariance
- use the learned eigenbasis
- build per-coordinate codebooks from PCA eigenvalues
- quantize normalized vectors and store norms separately

This is the SpectralQuant-style path.

### WHT-only

- use a fixed randomized Hadamard basis `D1 @ H @ D2`
- normalize each vector before rotation
- estimate empirical variance for each rotated coordinate from normalized calibration data
- build a separate codebook for each rotated coordinate using those empirical `sigma_i`
- quantize normalized vectors and store norms separately

This is the clean “WHT plus calibrated per-coordinate scalar quantization” branch.

### WHT PolarQuant-style

- use the same fixed randomized Hadamard basis
- normalize each vector before rotation
- use shared codebooks derived from the PolarQuant theory:
- `1-bit`: `±sqrt(2 / (pi d))`
- `2-bit`: fixed closed-form levels scaled by `1 / sqrt(d)`
- `3+ bit`: Lloyd-Max on `N(0, 1/d)`
- apply rotated-domain norm correction before inverse rotation
- store norms separately

This is the practical TurboQuant-style branch.

## Why The Earlier Runs Were Misleading

Earlier versions mixed together several inconsistencies:

- some runs compared proxy `Q=K` against real-`Q` causal
- some compared different models
- older “WHT” runs did not clearly separate `WHT-only` from `WHT PolarQuant-style`
- `WHT-only` originally used an incorrect flat `sigma = 1`
- `WHT-only` was later calibrated from raw vectors instead of normalized vectors

Those issues made the intermediate conclusions unstable.

The results below are from the final corrected version, where:

- `WHT-only` calibration is done on normalized data
- `WHT PolarQuant-style` uses the theoretical `1/d` scaling
- all three variants are produced by the same script in one run

## Commands

Proxy run:

```bash
.venv/bin/python experiments/codex_wht_vs_pca.py \
--device mps \
--mode main \
--model-name Qwen/Qwen2.5-1.5B-Instruct \
--short-name Qwen-1.5B
```

Real-`Q` causal run:

```bash
.venv/bin/python experiments/codex_wht_vs_pca.py \
--device mps \
--mode main \
--model-name Qwen/Qwen2.5-1.5B-Instruct \
--short-name Qwen-1.5B \
--evaluation-mode real_q_causal
```

## Raw Results

Proxy:

- `results/codex_wht_vs_pca/codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.json`
- `results/codex_wht_vs_pca/codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.md`

Real-`Q` causal:

- `results/codex_wht_vs_pca/codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.json`
- `results/codex_wht_vs_pca/codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.md`

## Final Results

### Proxy setup

Setup:

- `Q = K`
- non-causal attention
- model: `Qwen2.5-1.5B`

| Variant / Config | Attn Cos | KL | Key Cos | Value Cos |
|---|---:|---:|---:|---:|
| `PCA 1bit-K_6bit-V` | 0.7569 | 5.5865 | 0.9022 | 0.9695 |
| `WHT-only 1bit-K_6bit-V` | 0.8141 | 1.1236 | 0.8634 | 0.9996 |
| `WHT PolarQuant-style 1bit-K_6bit-V` | 0.9251 | 0.4098 | 0.8084 | 0.9996 |
| `PCA 1bit-K_4bit-V` | 0.6793 | 5.5865 | 0.9022 | 0.8709 |
| `WHT-only 1bit-K_4bit-V` | 0.8126 | 1.1236 | 0.8634 | 0.9956 |
| `WHT PolarQuant-style 1bit-K_4bit-V` | 0.9225 | 0.4098 | 0.8084 | 0.9955 |
| `PCA 2bit-K_6bit-V` | 0.7622 | 4.7245 | 0.9022 | 0.9695 |
| `WHT-only 2bit-K_6bit-V` | 0.9004 | 1.2201 | 0.9455 | 0.9996 |
| `WHT PolarQuant-style 2bit-K_6bit-V` | 0.9745 | 0.1114 | 0.9439 | 0.9996 |

### Real-`Q` causal setup

Setup:

- real captured `Q`
- causal attention
- model: `Qwen2.5-1.5B`

| Variant / Config | Attn Cos | KL | Key Cos | Value Cos |
|---|---:|---:|---:|---:|
| `PCA 1bit-K_6bit-V` | 0.5148 | 6.7447 | 0.9022 | 0.9695 |
| `WHT-only 1bit-K_6bit-V` | 0.7727 | 0.4972 | 0.8634 | 0.9996 |
| `WHT PolarQuant-style 1bit-K_6bit-V` | 0.6769 | 0.8792 | 0.8084 | 0.9996 |
| `PCA 1bit-K_4bit-V` | 0.4884 | 6.7447 | 0.9022 | 0.8709 |
| `WHT-only 1bit-K_4bit-V` | 0.7725 | 0.4972 | 0.8634 | 0.9956 |
| `WHT PolarQuant-style 1bit-K_4bit-V` | 0.6753 | 0.8792 | 0.8084 | 0.9955 |
| `PCA 2bit-K_6bit-V` | 0.5290 | 5.9388 | 0.9022 | 0.9695 |
| `WHT-only 2bit-K_6bit-V` | 0.8993 | 0.3452 | 0.9455 | 0.9996 |
| `WHT PolarQuant-style 2bit-K_6bit-V` | 0.8964 | 0.2739 | 0.9439 | 0.9996 |

## Interpretation

### 1. Is the learned PCA basis necessary?

No.

In both proxy and real-`Q` causal setups, the fixed WHT branches are better than PCA once the calibration is made internally consistent.

So the favorable `1-bit key` effect is not a PCA-only artifact.

### 2. Does WHT-only converge toward PolarQuant-style behavior?

Yes, much more than in the earlier buggy runs.

Once `WHT-only` estimates `sigma_i` from normalized rotated vectors, it moves into the same regime as the PolarQuant branch. This matches the theoretical expectation that the rotated normalized coordinates should have variance close to `1/d`.

There is still a difference:

- `WHT-only` uses empirical per-coordinate `sigma_i`
- `WHT PolarQuant-style` uses shared theoretical `1 / sqrt(d)`

But they are now clearly much closer than before.

### 3. Which branch is best?

Proxy:

- `WHT PolarQuant-style` > `WHT-only` > `PCA`

Real-`Q` causal:

- `WHT-only` > `WHT PolarQuant-style` > `PCA` on the 1-bit points
- `WHT-only` and `WHT PolarQuant-style` are essentially tied at `2bit-K_6bit-V`

So the final story is not “PolarQuant always dominates WHT-only.”

It is:

- both WHT branches are strong once calibrated consistently
- the basis change away from PCA matters
- the choice between empirical per-coordinate `sigma_i` and shared theoretical `1/sqrt(d)` is a second-order effect compared with the earlier calibration bugs

### 4. What do these results suggest?

They suggest that the strong performance is primarily coming from:

- normalizing the vectors
- using a randomized orthogonal WHT basis
- quantizing in that rotated normalized space

Once that is done correctly, both:

- empirical WHT calibration
- theoretical PolarQuant codebooks

work well.

## Bottom Line

The final cleaned-up experiment supports three conclusions:

1. The strong `1-bit key` result does not depend on the learned PCA basis.
2. The earlier contradictory WHT results were mostly caused by calibration inconsistencies.
3. After fixing those inconsistencies, WHT-based quantization is clearly competitive with or better than PCA, and the empirical-calibrated WHT and PolarQuant-style WHT branches land in the same general regime.
Loading