From 5eb4c4b0d3715323fac35a4f7fe9c5264ecdc82e Mon Sep 17 00:00:00 2001 From: Your Name Date: Thu, 9 Apr 2026 15:40:39 +0200 Subject: [PATCH 1/2] add Q=K ablation results --- README.md | 79 +++++ experiments/benchmark_utils.py | 329 ++++++++++++++++++ experiments/push_to_095.py | 180 +++++++--- .../push095_Qwen-0.5B-proxy-small.json | 74 ++++ .../push095_Qwen-0.5B-realq-small.json | 75 ++++ .../push095_Qwen-1.5B-proxy-small.json | 139 ++++++++ ...sh095_Qwen-1.5B-realq-noncausal-small.json | 227 ++++++++++++ ...ush095_Qwen-1.5B-realq-small-1bitcore.json | 75 ++++ .../push095_Qwen-1.5B-realq-small.json | 227 ++++++++++++ .../push095_TinyLlama-1.1B-proxy-small.json | 74 ++++ .../push095_TinyLlama-1.1B-realq-small.json | 75 ++++ .../qwen1p5b_1bit_claim_check_small.json | 279 +++++++++++++++ .../qwen1p5b_1bit_claim_check_small.md | 83 +++++ 13 files changed, 1876 insertions(+), 40 deletions(-) create mode 100644 experiments/benchmark_utils.py create mode 100644 results/push_095/push095_Qwen-0.5B-proxy-small.json create mode 100644 results/push_095/push095_Qwen-0.5B-realq-small.json create mode 100644 results/push_095/push095_Qwen-1.5B-proxy-small.json create mode 100644 results/push_095/push095_Qwen-1.5B-realq-noncausal-small.json create mode 100644 results/push_095/push095_Qwen-1.5B-realq-small-1bitcore.json create mode 100644 results/push_095/push095_Qwen-1.5B-realq-small.json create mode 100644 results/push_095/push095_TinyLlama-1.1B-proxy-small.json create mode 100644 results/push_095/push095_TinyLlama-1.1B-realq-small.json create mode 100644 results/push_095/qwen1p5b_1bit_claim_check_small.json create mode 100644 results/push_095/qwen1p5b_1bit_claim_check_small.md diff --git a/README.md b/README.md index d0e518a..8def8e6 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/experiments/benchmark_utils.py b/experiments/benchmark_utils.py new file mode 100644 index 0000000..007f197 --- /dev/null +++ b/experiments/benchmark_utils.py @@ -0,0 +1,329 @@ +import importlib +import inspect +import logging +import math +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple + +import torch +import torch.nn.functional as F + + +log = logging.getLogger(__name__) + + +COVERAGE_PRESETS: Dict[str, Dict[str, Any]] = { + "small": { + "n_calib": 16, + "n_eval": 8, + "seq_len": 512, + "layer_mode": "sampled", + "head_mode": "sampled", + "max_cached_vectors": 512, + }, + "medium": { + "n_calib": 64, + "n_eval": 32, + "seq_len": 1024, + "layer_mode": "sampled", + "head_mode": "sampled", + "max_cached_vectors": 1024, + }, + "large": { + "n_calib": 128, + "n_eval": 64, + "seq_len": 2048, + "layer_mode": "all", + "head_mode": "all", + "max_cached_vectors": 4096, + }, +} + + +def add_coverage_args(parser, *, default_preset: str = "medium") -> None: + parser.add_argument( + "--coverage-preset", + choices=sorted(COVERAGE_PRESETS), + default=default_preset, + help="Coverage preset controlling sample count, seq length, and layer/head breadth.", + ) + parser.add_argument("--n-calib", type=int, default=None) + parser.add_argument("--n-eval", type=int, default=None) + parser.add_argument("--seq-len", type=int, default=None) + parser.add_argument("--max-cached-vectors", type=int, default=None) + parser.add_argument( + "--layer-mode", + choices=["sampled", "all"], + default=None, + help="Use 3 representative layers or all layers.", + ) + parser.add_argument( + "--head-mode", + choices=["sampled", "all"], + default=None, + help="Use up to 4 KV heads or all KV heads.", + ) + + +def resolve_coverage_args(args) -> Dict[str, Any]: + preset = COVERAGE_PRESETS[args.coverage_preset].copy() + for key in ("n_calib", "n_eval", "seq_len", "max_cached_vectors", "layer_mode", "head_mode"): + value = getattr(args, key, None) + if value is not None: + preset[key] = value + preset["coverage_preset"] = args.coverage_preset + return preset + + +def select_layer_indices(n_layers: int, mode: str) -> List[int]: + if mode == "all": + return list(range(n_layers)) + return sorted(set([n_layers // 4, n_layers // 2, 3 * n_layers // 4])) + + +def select_head_indices(n_kv: int, mode: str) -> List[int]: + if mode == "all": + return list(range(n_kv)) + return list(range(min(n_kv, 4))) + + +def causal_attention_output(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, head_dim: int) -> torch.Tensor: + scale = 1.0 / math.sqrt(head_dim) + scores = Q.float() @ K.float().T * scale + q_len, k_len = scores.shape + if q_len == k_len: + mask = torch.triu( + torch.full((q_len, k_len), float("-inf"), device=scores.device, dtype=scores.dtype), + diagonal=1, + ) + scores = scores + mask + return F.softmax(scores, dim=-1) @ V.float() + + +def noncausal_attention_output(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, head_dim: int) -> torch.Tensor: + scale = 1.0 / math.sqrt(head_dim) + scores = Q.float() @ K.float().T * scale + return F.softmax(scores, dim=-1) @ V.float() + + +def _resolve_transformer_layers(model) -> Sequence[Any]: + candidates = [ + ("model", "layers"), + ("model", "decoder", "layers"), + ("transformer", "h"), + ("gpt_neox", "layers"), + ("layers",), + ] + for path in candidates: + cur = model + try: + for part in path: + cur = getattr(cur, part) + if len(cur): + return cur + except Exception: + continue + raise AttributeError("Could not locate transformer layers on model.") + + +def _resolve_attention_module(layer) -> Any: + for attr in ("self_attn", "attention", "attn"): + if hasattr(layer, attr): + return getattr(layer, attr) + raise AttributeError(f"Could not locate attention module on layer type {type(layer)!r}.") + + +def _shape_proj(proj: torch.Tensor, n_heads: int, head_dim: int) -> torch.Tensor: + batch, seq_len, _ = proj.shape + return proj.view(batch, seq_len, n_heads, head_dim).transpose(1, 2).contiguous() + + +def _capture_hidden_states(model, layer_indices: Iterable[int]): + layers = _resolve_transformer_layers(model) + captured: Dict[int, Dict[str, Any]] = {} + handles = [] + + for layer_idx in layer_indices: + attn = _resolve_attention_module(layers[layer_idx]) + + def _hook(module, args, kwargs, *, _layer_idx=layer_idx): + hidden_states = kwargs.get("hidden_states") + if hidden_states is None and args: + hidden_states = args[0] + if hidden_states is None: + return + captured[_layer_idx] = { + "hidden_states": hidden_states.detach(), + "args": args, + "kwargs": dict(kwargs), + } + + handles.append(attn.register_forward_pre_hook(_hook, with_kwargs=True)) + + return captured, handles + + +def _get_rotary_inputs(attn_module, hidden_states, n_kv, head_dim, capture): + if "position_embeddings" in capture["kwargs"]: + return capture["kwargs"]["position_embeddings"] + + position_ids = capture["kwargs"].get("position_ids") + if not hasattr(attn_module, "rotary_emb"): + raise ValueError("Attention module does not expose rotary embeddings.") + + value_states = _shape_proj(attn_module.v_proj(hidden_states), n_kv, head_dim) + rotary = attn_module.rotary_emb + try: + return rotary(value_states, position_ids) + except TypeError: + return rotary(value_states, position_ids=position_ids) + + +def _apply_rotary(attn_module, query_states, key_states, rotary_inputs): + if not isinstance(rotary_inputs, tuple) or len(rotary_inputs) != 2: + raise ValueError("Expected rotary inputs to be a (cos, sin) tuple.") + + mod = importlib.import_module(attn_module.__class__.__module__) + apply_fn = getattr(mod, "apply_rotary_pos_emb", None) + if apply_fn is None: + raise AttributeError(f"{mod.__name__} does not expose apply_rotary_pos_emb.") + + cos, sin = rotary_inputs + try: + return apply_fn(query_states, key_states, cos, sin, unsqueeze_dim=1) + except TypeError: + return apply_fn(query_states, key_states, cos, sin) + + +def compute_true_queries_for_layer(model, layer_idx: int, capture: Dict[str, Any], head_dim: int) -> torch.Tensor: + layers = _resolve_transformer_layers(model) + attn = _resolve_attention_module(layers[layer_idx]) + hidden_states = capture["hidden_states"].to(attn.q_proj.weight.device, dtype=attn.q_proj.weight.dtype) + + cfg = model.config + n_heads = int(getattr(cfg, "num_attention_heads")) + n_kv = int(getattr(cfg, "num_key_value_heads", n_heads)) + + query_states = _shape_proj(attn.q_proj(hidden_states), n_heads, head_dim) + key_states = _shape_proj(attn.k_proj(hidden_states), n_kv, head_dim) + + try: + rotary_inputs = _get_rotary_inputs(attn, hidden_states, n_kv, head_dim, capture) + query_states, _ = _apply_rotary(attn, query_states, key_states, rotary_inputs) + except Exception as exc: + log.warning("Layer %d query capture fell back to pre-RoPE projections: %s", layer_idx, exc) + + return query_states.detach().float().cpu() + + +def map_kv_head_to_query_head(kv_head_idx: int, n_heads: int, n_kv: int) -> int: + if n_heads <= n_kv: + return min(kv_head_idx, n_heads - 1) + q_per_kv = max(n_heads // n_kv, 1) + return min(kv_head_idx * q_per_kv, n_heads - 1) + + +def collect_eval_qkv( + model, + tokenizer, + device: str, + *, + n_eval: int, + seq_len: int, + head_dim: int, + layer_indices: Sequence[int], + head_indices: Sequence[int], + extract_kv_layer, + dataset_name: str = "wikitext", + dataset_config: str = "wikitext-103-raw-v1", + split: str = "train", + split_start: int = 0, +) -> Dict[str, Any]: + from datasets import load_dataset + + cfg = model.config + n_heads = int(getattr(cfg, "num_attention_heads")) + n_kv = int(getattr(cfg, "num_key_value_heads", n_heads)) + dataset_kwargs = { + "path": dataset_name, + "name": dataset_config, + "split": f"{split}[{split_start}:{split_start + n_eval * 5}]", + } + try: + ds = load_dataset(**dataset_kwargs) + except Exception as exc: + log.warning("Dataset load failed, retrying from local cache only: %s", exc) + ds = load_dataset(**dataset_kwargs, download_mode="reuse_cache_if_exists") + eval_queries = {(l, h): [] for l in layer_indices for h in head_indices} + eval_keys = {(l, h): [] for l in layer_indices for h in head_indices} + eval_vals = {(l, h): [] for l in layer_indices for h in head_indices} + + captured, handles = _capture_hidden_states(model, layer_indices) + n_sequences = 0 + try: + for item in ds: + text = item.get("text", "") + if len(text.strip()) < 100: + continue + + enc = tokenizer(text, return_tensors="pt", max_length=seq_len, truncation=True).to(device) + if enc["input_ids"].shape[1] < 32: + continue + + captured.clear() + with torch.no_grad(): + out = model(**enc, use_cache=True) + kv = out.past_key_values + + layer_queries: Dict[int, torch.Tensor] = {} + for layer_idx in layer_indices: + capture = captured.get(layer_idx) + if capture is None: + continue + layer_queries[layer_idx] = compute_true_queries_for_layer(model, layer_idx, capture, head_dim) + + for layer_idx in layer_indices: + try: + k_l, v_l = extract_kv_layer(kv, layer_idx) + k_l = k_l.float().cpu() + v_l = v_l.float().cpu() + except Exception: + continue + + q_l = layer_queries.get(layer_idx) + if q_l is None: + continue + + for kv_head_idx in head_indices: + if kv_head_idx >= k_l.shape[1]: + continue + query_head_idx = map_kv_head_to_query_head(kv_head_idx, n_heads=n_heads, n_kv=n_kv) + if query_head_idx >= q_l.shape[1]: + continue + eval_queries[(layer_idx, kv_head_idx)].append(q_l[0, query_head_idx]) + eval_keys[(layer_idx, kv_head_idx)].append(k_l[0, kv_head_idx]) + eval_vals[(layer_idx, kv_head_idx)].append(v_l[0, kv_head_idx]) + + n_sequences += 1 + if n_sequences >= n_eval: + break + finally: + for handle in handles: + handle.remove() + + return { + "queries": eval_queries, + "keys": eval_keys, + "values": eval_vals, + "n_sequences": n_sequences, + "query_mode": "true_model_queries", + "query_head_mapping": "first_query_head_per_kv_group", + "attention_mode": "causal_self_attention", + "seq_len": seq_len, + } + + +def build_benchmark_note(key_bits: Any) -> str: + return ( + f"{key_bits}-bit keys here means per-dimension scalar quantization in the rotated spectral basis, " + "plus a separately stored vector norm. It is not a raw sign-only key in the original basis." + ) diff --git a/experiments/push_to_095.py b/experiments/push_to_095.py index a678699..ca054ad 100644 --- a/experiments/push_to_095.py +++ b/experiments/push_to_095.py @@ -29,6 +29,17 @@ sys.path.insert(0, str(PROJECT_ROOT / "src")) sys.path.insert(0, str(PROJECT_ROOT / "baseline" / "turboquant_cutile")) +from benchmark_utils import ( + add_coverage_args, + build_benchmark_note, + causal_attention_output, + collect_eval_qkv, + noncausal_attention_output, + resolve_coverage_args, + select_head_indices, + select_layer_indices, +) + log = logging.getLogger("push095") logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s - %(message)s") @@ -44,6 +55,8 @@ N_CALIB = 64 N_EVAL = 32 +SEQ_LEN = 512 +MAX_CACHED_VECTORS = 512 def save_result(filename, data): @@ -109,9 +122,37 @@ def load_model_tokenizer(model_name, device): needs_token = any(x in model_name for x in ["llama", "Llama", "gemma", "Gemma"]) token = HF_TOKEN if needs_token else None log.info("Loading %s ...", model_name) - tokenizer = AutoTokenizer.from_pretrained(model_name, token=token) + try: + tokenizer = AutoTokenizer.from_pretrained(model_name, token=token) + except Exception as exc: + log.warning("Tokenizer load failed, retrying from local cache only: %s", exc) + tokenizer = AutoTokenizer.from_pretrained(model_name, token=token, local_files_only=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map=device, token=token) + model_kwargs = {"token": token} + if device == "cpu": + model_kwargs["torch_dtype"] = torch.float32 + try: + model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) + except Exception as exc: + log.warning("Model load failed, retrying from local cache only: %s", exc) + model = AutoModelForCausalLM.from_pretrained(model_name, local_files_only=True, **model_kwargs) + model.to(device) + elif device == "mps": + model_kwargs["torch_dtype"] = torch.float16 + try: + model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) + except Exception as exc: + log.warning("Model load failed, retrying from local cache only: %s", exc) + model = AutoModelForCausalLM.from_pretrained(model_name, local_files_only=True, **model_kwargs) + model.to(device) + else: + model_kwargs["torch_dtype"] = torch.float16 + model_kwargs["device_map"] = device + try: + model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) + except Exception as exc: + log.warning("Model load failed, retrying from local cache only: %s", exc) + model = AutoModelForCausalLM.from_pretrained(model_name, local_files_only=True, **model_kwargs) model.eval() cfg = model.config; n_layers = cfg.num_hidden_layers n_kv = getattr(cfg, "num_key_value_heads", cfg.num_attention_heads) @@ -223,13 +264,12 @@ def decompress(self, compressed, VT): qr = compressed["qr"].float(); norms = compressed["norms"].float() return (qr @ VT.to(qr.device)) * norms.unsqueeze(-1) - def evaluate(self, Q, K, V, fp16_output): + def evaluate(self, Q, K, V, fp16_output, *, attention_fn): ck = self.compress_keys(K); cv = self.compress_values(V) K_hat = self.decompress(ck, self.VkT); V_hat = self.decompress(cv, self.VvT) key_cos = F.cosine_similarity(K.float(), K_hat.float(), dim=-1).mean().item() val_cos = F.cosine_similarity(V.float(), V_hat.float(), dim=-1).mean().item() - scale = 1.0 / math.sqrt(self.hd) - output = F.softmax(Q.float() @ K_hat.T * scale, dim=-1) @ V_hat + output = attention_fn(Q, K_hat, V_hat, self.hd) attn_cos = F.cosine_similarity(fp16_output.float(), output.float(), dim=-1).mean().item() return {"attn_cos_sim": attn_cos, "key_cos_sim": key_cos, "val_cos_sim": val_cos, "key_bits": self.key_total_bits, "val_bits": self.val_total_bits, @@ -237,41 +277,55 @@ def evaluate(self, Q, K, V, fp16_output): "total_compress": (2*self.hd*16) / (self.key_total_bits + self.val_total_bits)} -def run_model(model_name, short_name, device): +def run_model(model_name, short_name, device, coverage, *, query_mode: str, attention_mode: str): log.info("=" * 70) log.info("MODEL: %s", model_name) model, tokenizer, n_layers, n_kv, hd = load_model_tokenizer(model_name, device) - eigen_keys, eigen_vals = calibrate(model, tokenizer, N_CALIB, device, n_layers, n_kv, hd) + n_calib = coverage["n_calib"] + n_eval = coverage["n_eval"] + seq_len = coverage["seq_len"] + max_cached_vectors = coverage["max_cached_vectors"] + layer_indices = select_layer_indices(n_layers, coverage["layer_mode"]) + head_indices = select_head_indices(n_kv, coverage["head_mode"]) + + eigen_keys, eigen_vals = calibrate(model, tokenizer, n_calib, device, n_layers, n_kv, hd) deff_k = np.mean([eigen_keys[(l,h)]["d_eff"] for l in range(n_layers) for h in range(n_kv)]) deff_v = np.mean([eigen_vals[(l,h)]["d_eff"] for l in range(n_layers) for h in range(n_kv)]) log.info("d_eff: keys=%.1f, values=%.1f", deff_k, deff_v) - layer_indices = sorted(set([n_layers//4, n_layers//2, 3*n_layers//4])) - head_indices = list(range(min(n_kv, 4))) - - from datasets import load_dataset - ds = load_dataset("wikitext", "wikitext-103-raw-v1", split=f"train[{N_CALIB*5}:{(N_CALIB+N_EVAL)*5}]") - eval_keys = {(l,h): [] for l in layer_indices for h in head_indices} - eval_vals = {(l,h): [] for l in layer_indices for h in head_indices} - nd = 0 - for item in ds: - text = item.get("text","") - if len(text.strip()) < 100: continue - enc = tokenizer(text, return_tensors="pt", max_length=512, truncation=True).to(device) - if enc["input_ids"].shape[1] < 32: continue - with torch.no_grad(): - out = model(**enc, use_cache=True); kv = out.past_key_values - for l in layer_indices: - try: k_l, v_l = extract_kv_layer(kv, l) - except: continue - k_l = k_l.float().cpu(); v_l = v_l.float().cpu() - for h in head_indices: - if h < k_l.shape[1]: - eval_keys[(l,h)].append(k_l[0,h]); eval_vals[(l,h)].append(v_l[0,h]) - nd += 1 - if nd >= N_EVAL: break + eval_bundle = collect_eval_qkv( + model, + tokenizer, + device, + n_eval=n_eval, + seq_len=seq_len, + head_dim=hd, + layer_indices=layer_indices, + head_indices=head_indices, + extract_kv_layer=extract_kv_layer, + split_start=n_calib * 5, + ) + eval_queries = eval_bundle["queries"] + eval_keys = eval_bundle["keys"] + eval_vals = eval_bundle["values"] + log.info("Collected eval data from %d sequences", eval_bundle["n_sequences"]) + + if query_mode == "proxy_qeqk": + eval_queries = eval_keys + query_mode_label = "proxy_keys_as_queries" + query_head_mapping = None + else: + query_mode_label = eval_bundle["query_mode"] + query_head_mapping = eval_bundle["query_head_mapping"] + + if attention_mode == "causal": + attention_fn = causal_attention_output + attention_mode_label = "causal_self_attention" + else: + attention_fn = noncausal_attention_output + attention_mode_label = "noncausal_self_attention_proxy" # Define configs: (name, key_bits_per_dim, val_mode, val_param) configs = [ @@ -303,15 +357,15 @@ def run_model(model_name, short_name, device): for l in layer_indices: for h in head_indices: + qq = eval_queries.get((l,h),[]) kk = eval_keys.get((l,h),[]); vv = eval_vals.get((l,h),[]) - if not kk or not vv: continue - K_all = torch.cat(kk, dim=0)[:512].to(device).float() - V_all = torch.cat(vv, dim=0)[:512].to(device).float() - Q_all = K_all.clone() + if not qq or not kk or not vv: continue + Q_all = torch.cat(qq, dim=0)[:max_cached_vectors].to(device).float() + K_all = torch.cat(kk, dim=0)[:max_cached_vectors].to(device).float() + V_all = torch.cat(vv, dim=0)[:max_cached_vectors].to(device).float() if K_all.shape[0] < 32: continue - scale = 1.0 / math.sqrt(hd) - fp16_output = F.softmax(Q_all @ K_all.T * scale, dim=-1) @ V_all + fp16_output = attention_fn(Q_all, K_all, V_all, hd) v_ev = eigen_vals[(l,h)]["ev"].cpu().numpy() if val_mode == "uniform": @@ -324,7 +378,7 @@ def run_model(model_name, short_name, device): eigen_vals[(l,h)]["evec"].to(device), eigen_vals[(l,h)]["ev"].to(device), key_bpd, val_alloc, hd) - m = engine.evaluate(Q_all, K_all, V_all, fp16_output) + m = engine.evaluate(Q_all, K_all, V_all, fp16_output, attention_fn=attention_fn) metrics_list.append(m) if metrics_list: @@ -350,6 +404,23 @@ def run_model(model_name, short_name, device): "head_dim": hd, "n_layers": n_layers, "n_kv_heads": n_kv, "mean_d_eff_keys": float(deff_k), "mean_d_eff_values": float(deff_v), "configs": all_results, + "benchmark_note": build_benchmark_note(1), + "query_capture": { + "mode": query_mode_label, + "head_mapping": query_head_mapping, + "attention_mode": attention_mode_label, + }, + "coverage": { + "preset": coverage["coverage_preset"], + "n_calib": n_calib, + "n_eval": n_eval, + "seq_len": seq_len, + "max_cached_vectors": max_cached_vectors, + "layer_mode": coverage["layer_mode"], + "head_mode": coverage["head_mode"], + "layers_evaluated": layer_indices, + "heads_evaluated": head_indices, + }, } save_result(f"push095_{short_name}.json", result) @@ -361,12 +432,41 @@ def run_model(model_name, short_name, device): def main(): parser = argparse.ArgumentParser() parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + parser.add_argument("--model-name", default=None, help="Optional single model to run.") + parser.add_argument("--short-name", default=None, help="Required when --model-name is set.") + parser.add_argument( + "--query-mode", + choices=["real_q", "proxy_qeqk"], + default="real_q", + help="Use captured model queries or the Q=K proxy.", + ) + parser.add_argument( + "--attention-mode", + choices=["causal", "noncausal"], + default="causal", + help="Evaluate attention with a causal mask or a non-causal proxy.", + ) + add_coverage_args(parser, default_preset="medium") args = parser.parse_args() + coverage = resolve_coverage_args(args) + + models = MODELS + if args.model_name: + if not args.short_name: + parser.error("--short-name is required when --model-name is set") + models = [(args.model_name, args.short_name)] all_results = {}; t0 = time.time() - for model_name, short_name in MODELS: + for model_name, short_name in models: try: - result = run_model(model_name, short_name, args.device) + result = run_model( + model_name, + short_name, + args.device, + coverage, + query_mode=args.query_mode, + attention_mode=args.attention_mode, + ) all_results[short_name] = result except Exception as e: log.error("FAILED on %s: %s", model_name, e) diff --git a/results/push_095/push095_Qwen-0.5B-proxy-small.json b/results/push_095/push095_Qwen-0.5B-proxy-small.json new file mode 100644 index 0000000..bacb409 --- /dev/null +++ b/results/push_095/push095_Qwen-0.5B-proxy-small.json @@ -0,0 +1,74 @@ +{ + "model": "Qwen/Qwen2.5-0.5B-Instruct", + "short_name": "Qwen-0.5B-small", + "coverage": { + "n_calib": 8, + "n_eval": 4, + "seq_len": 512, + "layer_mode": "sampled", + "head_mode": "sampled", + "max_cached_vectors": 256, + "coverage_preset": "small" + }, + "query_capture": { + "mode": "proxy_keys_as_queries", + "attention_mode": "noncausal_self_attention_proxy" + }, + "configs": { + "1bit-K_3bit-V": { + "attn_cos_sim": 0.7519614895184835, + "attn_cos_sim_std": 0.06812770604500246, + "key_cos_sim": 0.9039093752702078, + "val_cos_sim": 0.8488263885180155, + "key_bits": 80, + "val_bits": 208, + "total_bits": 288, + "total_compress": 7.111111111111111, + "n_evals": 6 + }, + "1bit-K_4bit-V": { + "attn_cos_sim": 0.8217491110165914, + "attn_cos_sim_std": 0.07753627443686101, + "key_cos_sim": 0.9039093752702078, + "val_cos_sim": 0.9384434620539347, + "key_bits": 80, + "val_bits": 272, + "total_bits": 352, + "total_compress": 5.818181818181818, + "n_evals": 6 + }, + "1bit-K_5bit-V": { + "attn_cos_sim": 0.8587081730365753, + "attn_cos_sim_std": 0.08041484500549338, + "key_cos_sim": 0.9039093752702078, + "val_cos_sim": 0.979455441236496, + "key_bits": 80, + "val_bits": 336, + "total_bits": 416, + "total_compress": 4.923076923076923, + "n_evals": 6 + }, + "1bit-K_6bit-V": { + "attn_cos_sim": 0.8676428099473318, + "attn_cos_sim_std": 0.0821010315119471, + "key_cos_sim": 0.9039093752702078, + "val_cos_sim": 0.9915579358736674, + "key_bits": 80, + "val_bits": 400, + "total_bits": 480, + "total_compress": 4.266666666666667, + "n_evals": 6 + }, + "uniform-3bit": { + "attn_cos_sim": 0.7794286708037058, + "attn_cos_sim_std": 0.055274255601519995, + "key_cos_sim": 0.9044231275717417, + "val_cos_sim": 0.8488263885180155, + "key_bits": 208, + "val_bits": 208, + "total_bits": 416, + "total_compress": 4.923076923076923, + "n_evals": 6 + } + } +} \ No newline at end of file diff --git a/results/push_095/push095_Qwen-0.5B-realq-small.json b/results/push_095/push095_Qwen-0.5B-realq-small.json new file mode 100644 index 0000000..207ff41 --- /dev/null +++ b/results/push_095/push095_Qwen-0.5B-realq-small.json @@ -0,0 +1,75 @@ +{ + "model": "Qwen/Qwen2.5-0.5B-Instruct", + "short_name": "Qwen-0.5B-small", + "coverage": { + "n_calib": 8, + "n_eval": 4, + "seq_len": 512, + "layer_mode": "sampled", + "head_mode": "sampled", + "max_cached_vectors": 256, + "coverage_preset": "small" + }, + "query_capture": { + "mode": "true_model_queries", + "head_mapping": "first_query_head_per_kv_group", + "attention_mode": "causal_self_attention" + }, + "configs": { + "1bit-K_3bit-V": { + "attn_cos_sim": 0.5295048902432123, + "attn_cos_sim_std": 0.046038966891283103, + "key_cos_sim": 0.9039093752702078, + "val_cos_sim": 0.8488263885180155, + "key_bits": 80, + "val_bits": 208, + "total_bits": 288, + "total_compress": 7.111111111111111, + "n_evals": 6 + }, + "1bit-K_4bit-V": { + "attn_cos_sim": 0.5828911165396372, + "attn_cos_sim_std": 0.05879976582914788, + "key_cos_sim": 0.9039093752702078, + "val_cos_sim": 0.9384434620539347, + "key_bits": 80, + "val_bits": 272, + "total_bits": 352, + "total_compress": 5.818181818181818, + "n_evals": 6 + }, + "1bit-K_5bit-V": { + "attn_cos_sim": 0.6017550925413767, + "attn_cos_sim_std": 0.06387837533517161, + "key_cos_sim": 0.9039093752702078, + "val_cos_sim": 0.979455441236496, + "key_bits": 80, + "val_bits": 336, + "total_bits": 416, + "total_compress": 4.923076923076923, + "n_evals": 6 + }, + "1bit-K_6bit-V": { + "attn_cos_sim": 0.6070036987463633, + "attn_cos_sim_std": 0.06712388329158356, + "key_cos_sim": 0.9039093752702078, + "val_cos_sim": 0.9915579358736674, + "key_bits": 80, + "val_bits": 400, + "total_bits": 480, + "total_compress": 4.266666666666667, + "n_evals": 6 + }, + "uniform-3bit": { + "attn_cos_sim": 0.5817577640215555, + "attn_cos_sim_std": 0.04111165374240058, + "key_cos_sim": 0.9044231275717417, + "val_cos_sim": 0.8488263885180155, + "key_bits": 208, + "val_bits": 208, + "total_bits": 416, + "total_compress": 4.923076923076923, + "n_evals": 6 + } + } +} \ No newline at end of file diff --git a/results/push_095/push095_Qwen-1.5B-proxy-small.json b/results/push_095/push095_Qwen-1.5B-proxy-small.json new file mode 100644 index 0000000..38a4444 --- /dev/null +++ b/results/push_095/push095_Qwen-1.5B-proxy-small.json @@ -0,0 +1,139 @@ +{ + "model": "Qwen/Qwen2.5-1.5B-Instruct", + "short_name": "Qwen-1.5B-proxy-small", + "coverage": { + "n_calib": 8, + "n_eval": 4, + "seq_len": 512, + "layer_mode": "sampled", + "head_mode": "sampled", + "max_cached_vectors": 256, + "coverage_preset": "small" + }, + "query_capture": { + "mode": "proxy_keys_as_queries", + "attention_mode": "noncausal_self_attention_proxy" + }, + "configs": { + "1bit-K_3bit-V": { + "config": "1bit-K_3bit-V", + "attn_cos_sim": 0.649636780222257, + "attn_cos_sim_std": 0.2728981431678952, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8254336913426717, + "key_bits": 144, + "val_bits": 400, + "total_bits": 544, + "total_compress": 7.529411764705882, + "n_evals": 6 + }, + "1bit-K_4bit-V": { + "config": "1bit-K_4bit-V", + "attn_cos_sim": 0.6717648692429066, + "attn_cos_sim_std": 0.2813665016012008, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8687977194786072, + "key_bits": 144, + "val_bits": 528, + "total_bits": 672, + "total_compress": 6.095238095238095, + "n_evals": 6 + }, + "1bit-K_5bit-V": { + "config": "1bit-K_5bit-V", + "attn_cos_sim": 0.7147521277268728, + "attn_cos_sim_std": 0.2979772684157822, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9365533093611399, + "key_bits": 144, + "val_bits": 656, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "1bit-K_6bit-V": { + "config": "1bit-K_6bit-V", + "attn_cos_sim": 0.7450543741385142, + "attn_cos_sim_std": 0.3066504123342587, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9689570963382721, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "1bit-K_optV-400": { + "config": "1bit-K_optV-400", + "attn_cos_sim": 0.5912532961616913, + "attn_cos_sim_std": 0.2508941894516606, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.7956738471984863, + "key_bits": 144, + "val_bits": 416, + "total_bits": 560, + "total_compress": 7.314285714285714, + "n_evals": 6 + }, + "1bit-K_optV-500": { + "config": "1bit-K_optV-500", + "attn_cos_sim": 0.6651112387577692, + "attn_cos_sim_std": 0.2759489913732797, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8773498833179474, + "key_bits": 144, + "val_bits": 516, + "total_bits": 660, + "total_compress": 6.206060606060606, + "n_evals": 6 + }, + "1bit-K_optV-600": { + "config": "1bit-K_optV-600", + "attn_cos_sim": 0.7170996777713299, + "attn_cos_sim_std": 0.2980135224683886, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9368999799092611, + "key_bits": 144, + "val_bits": 616, + "total_bits": 760, + "total_compress": 5.389473684210526, + "n_evals": 6 + }, + "1bit-K_optV-700": { + "config": "1bit-K_optV-700", + "attn_cos_sim": 0.7395737692713737, + "attn_cos_sim_std": 0.3046034646473958, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9626753330230713, + "key_bits": 144, + "val_bits": 716, + "total_bits": 860, + "total_compress": 4.7627906976744185, + "n_evals": 6 + }, + "1bit-K_optV-800": { + "config": "1bit-K_optV-800", + "attn_cos_sim": 0.7450543741385142, + "attn_cos_sim_std": 0.3066504123342587, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9689570963382721, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "uniform-3bit": { + "config": "uniform-3bit", + "attn_cos_sim": 0.6673308735092481, + "attn_cos_sim_std": 0.2781717269197852, + "key_cos_sim": 0.9064196844895681, + "val_cos_sim": 0.8254336913426717, + "key_bits": 400, + "val_bits": 400, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + } + } +} \ No newline at end of file diff --git a/results/push_095/push095_Qwen-1.5B-realq-noncausal-small.json b/results/push_095/push095_Qwen-1.5B-realq-noncausal-small.json new file mode 100644 index 0000000..b1c898d --- /dev/null +++ b/results/push_095/push095_Qwen-1.5B-realq-noncausal-small.json @@ -0,0 +1,227 @@ +{ + "model": "Qwen/Qwen2.5-1.5B-Instruct", + "short_name": "Qwen-1.5B-realq-noncausal-small", + "head_dim": 128, + "n_layers": 28, + "n_kv_heads": 2, + "mean_d_eff_keys": 3.8392857142857144, + "mean_d_eff_values": 35.05357142857143, + "configs": { + "1bit-K_3bit-V": { + "config": "1bit-K_3bit-V", + "attn_cos_sim": 0.36049971108635265, + "attn_cos_sim_std": 0.23867883449471655, + "key_cos_sim": 0.9062796235084534, + "val_cos_sim": 0.825449655453364, + "key_bits": 144, + "val_bits": 400, + "total_bits": 544, + "total_compress": 7.529411764705882, + "n_evals": 6 + }, + "1bit-K_4bit-V": { + "config": "1bit-K_4bit-V", + "attn_cos_sim": 0.3948807418346405, + "attn_cos_sim_std": 0.22874032956813606, + "key_cos_sim": 0.9062796235084534, + "val_cos_sim": 0.8687503139177958, + "key_bits": 144, + "val_bits": 528, + "total_bits": 672, + "total_compress": 6.095238095238095, + "n_evals": 6 + }, + "1bit-K_5bit-V": { + "config": "1bit-K_5bit-V", + "attn_cos_sim": 0.4024066738784313, + "attn_cos_sim_std": 0.23151242013206327, + "key_cos_sim": 0.9062796235084534, + "val_cos_sim": 0.9365509152412415, + "key_bits": 144, + "val_bits": 656, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "1bit-K_6bit-V": { + "config": "1bit-K_6bit-V", + "attn_cos_sim": 0.41242237637440365, + "attn_cos_sim_std": 0.23545930203389223, + "key_cos_sim": 0.9062796235084534, + "val_cos_sim": 0.9689375261465708, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "2bit-K_3bit-V": { + "config": "2bit-K_3bit-V", + "attn_cos_sim": 0.36948591594894725, + "attn_cos_sim_std": 0.23901105726822372, + "key_cos_sim": 0.9063107967376709, + "val_cos_sim": 0.825449655453364, + "key_bits": 272, + "val_bits": 400, + "total_bits": 672, + "total_compress": 6.095238095238095, + "n_evals": 6 + }, + "2bit-K_4bit-V": { + "config": "2bit-K_4bit-V", + "attn_cos_sim": 0.4043494140108426, + "attn_cos_sim_std": 0.22904510409405765, + "key_cos_sim": 0.9063107967376709, + "val_cos_sim": 0.8687503139177958, + "key_bits": 272, + "val_bits": 528, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "2bit-K_5bit-V": { + "config": "2bit-K_5bit-V", + "attn_cos_sim": 0.41210776567459106, + "attn_cos_sim_std": 0.2328154502237668, + "key_cos_sim": 0.9063107967376709, + "val_cos_sim": 0.9365509152412415, + "key_bits": 272, + "val_bits": 656, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "1bit-K_optV-400": { + "config": "1bit-K_optV-400", + "attn_cos_sim": 0.3446629196405411, + "attn_cos_sim_std": 0.18278215712698637, + "key_cos_sim": 0.9062796235084534, + "val_cos_sim": 0.7953943014144897, + "key_bits": 144, + "val_bits": 416, + "total_bits": 560, + "total_compress": 7.314285714285714, + "n_evals": 6 + }, + "1bit-K_optV-500": { + "config": "1bit-K_optV-500", + "attn_cos_sim": 0.3886757865548134, + "attn_cos_sim_std": 0.20960839569248663, + "key_cos_sim": 0.9062796235084534, + "val_cos_sim": 0.8773297071456909, + "key_bits": 144, + "val_bits": 516, + "total_bits": 660, + "total_compress": 6.206060606060606, + "n_evals": 6 + }, + "1bit-K_optV-600": { + "config": "1bit-K_optV-600", + "attn_cos_sim": 0.40516414990027744, + "attn_cos_sim_std": 0.2305832759974119, + "key_cos_sim": 0.9062796235084534, + "val_cos_sim": 0.9368569254875183, + "key_bits": 144, + "val_bits": 616, + "total_bits": 760, + "total_compress": 5.389473684210526, + "n_evals": 6 + }, + "1bit-K_optV-700": { + "config": "1bit-K_optV-700", + "attn_cos_sim": 0.407974520077308, + "attn_cos_sim_std": 0.23506802378146752, + "key_cos_sim": 0.9062796235084534, + "val_cos_sim": 0.9626708130041758, + "key_bits": 144, + "val_bits": 716, + "total_bits": 860, + "total_compress": 4.7627906976744185, + "n_evals": 6 + }, + "1bit-K_optV-800": { + "config": "1bit-K_optV-800", + "attn_cos_sim": 0.41242237637440365, + "attn_cos_sim_std": 0.23545930203389223, + "key_cos_sim": 0.9062796235084534, + "val_cos_sim": 0.9689375261465708, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "2bit-K_optV-500": { + "config": "2bit-K_optV-500", + "attn_cos_sim": 0.39863120516141254, + "attn_cos_sim_std": 0.21093295517205174, + "key_cos_sim": 0.9063107967376709, + "val_cos_sim": 0.8773297071456909, + "key_bits": 272, + "val_bits": 516, + "total_bits": 788, + "total_compress": 5.197969543147208, + "n_evals": 6 + }, + "2bit-K_optV-600": { + "config": "2bit-K_optV-600", + "attn_cos_sim": 0.41510626549522084, + "attn_cos_sim_std": 0.23167674501789354, + "key_cos_sim": 0.9063107967376709, + "val_cos_sim": 0.9368569254875183, + "key_bits": 272, + "val_bits": 616, + "total_bits": 888, + "total_compress": 4.612612612612613, + "n_evals": 6 + }, + "uniform-3bit": { + "config": "uniform-3bit", + "attn_cos_sim": 0.3900533703466256, + "attn_cos_sim_std": 0.23825334122423134, + "key_cos_sim": 0.9064055383205414, + "val_cos_sim": 0.825449655453364, + "key_bits": 400, + "val_bits": 400, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "uniform-4bit": { + "config": "uniform-4bit", + "attn_cos_sim": 0.48846515516440075, + "attn_cos_sim_std": 0.22806892138519574, + "key_cos_sim": 0.9068096180756887, + "val_cos_sim": 0.8687503139177958, + "key_bits": 528, + "val_bits": 528, + "total_bits": 1056, + "total_compress": 3.878787878787879, + "n_evals": 6 + } + }, + "benchmark_note": "1-bit keys here means per-dimension scalar quantization in the rotated spectral basis, plus a separately stored vector norm. It is not a raw sign-only key in the original basis.", + "query_capture": { + "mode": "true_model_queries", + "head_mapping": "first_query_head_per_kv_group", + "attention_mode": "noncausal_self_attention_proxy" + }, + "coverage": { + "preset": "small", + "n_calib": 8, + "n_eval": 4, + "seq_len": 512, + "max_cached_vectors": 256, + "layer_mode": "sampled", + "head_mode": "sampled", + "layers_evaluated": [ + 7, + 14, + 21 + ], + "heads_evaluated": [ + 0, + 1 + ] + } +} \ No newline at end of file diff --git a/results/push_095/push095_Qwen-1.5B-realq-small-1bitcore.json b/results/push_095/push095_Qwen-1.5B-realq-small-1bitcore.json new file mode 100644 index 0000000..ff611d3 --- /dev/null +++ b/results/push_095/push095_Qwen-1.5B-realq-small-1bitcore.json @@ -0,0 +1,75 @@ +{ + "model": "Qwen/Qwen2.5-1.5B-Instruct", + "short_name": "Qwen-1.5B-realq-small-1bitcore", + "coverage": { + "n_calib": 8, + "n_eval": 4, + "seq_len": 512, + "layer_mode": "sampled", + "head_mode": "sampled", + "max_cached_vectors": 256, + "coverage_preset": "small" + }, + "query_capture": { + "mode": "true_model_queries", + "head_mapping": "first_query_head_per_kv_group", + "attention_mode": "causal_self_attention" + }, + "configs": { + "1bit-K_3bit-V": { + "attn_cos_sim": 0.4363829667369525, + "attn_cos_sim_std": 0.19861484805362176, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8254336913426717, + "key_bits": 144, + "val_bits": 400, + "total_bits": 544, + "total_compress": 7.529411764705882, + "n_evals": 6 + }, + "1bit-K_4bit-V": { + "attn_cos_sim": 0.47968902190526325, + "attn_cos_sim_std": 0.18598941169446778, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8687977194786072, + "key_bits": 144, + "val_bits": 528, + "total_bits": 672, + "total_compress": 6.095238095238095, + "n_evals": 6 + }, + "1bit-K_5bit-V": { + "attn_cos_sim": 0.5087136526902517, + "attn_cos_sim_std": 0.2049756014656585, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9365533093611399, + "key_bits": 144, + "val_bits": 656, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "1bit-K_6bit-V": { + "attn_cos_sim": 0.5249765266974767, + "attn_cos_sim_std": 0.19193184629905366, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9689570963382721, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "uniform-3bit": { + "attn_cos_sim": 0.45800502101580304, + "attn_cos_sim_std": 0.19474016746381387, + "key_cos_sim": 0.9064196844895681, + "val_cos_sim": 0.8254336913426717, + "key_bits": 400, + "val_bits": 400, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + } + } +} \ No newline at end of file diff --git a/results/push_095/push095_Qwen-1.5B-realq-small.json b/results/push_095/push095_Qwen-1.5B-realq-small.json new file mode 100644 index 0000000..b4bee3f --- /dev/null +++ b/results/push_095/push095_Qwen-1.5B-realq-small.json @@ -0,0 +1,227 @@ +{ + "model": "Qwen/Qwen2.5-1.5B-Instruct", + "short_name": "Qwen-1.5B-realq-small", + "head_dim": 128, + "n_layers": 28, + "n_kv_heads": 2, + "mean_d_eff_keys": 3.8392857142857144, + "mean_d_eff_values": 35.07142857142857, + "configs": { + "1bit-K_3bit-V": { + "config": "1bit-K_3bit-V", + "attn_cos_sim": 0.4363829667369525, + "attn_cos_sim_std": 0.19861484805362176, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8254336913426717, + "key_bits": 144, + "val_bits": 400, + "total_bits": 544, + "total_compress": 7.529411764705882, + "n_evals": 6 + }, + "1bit-K_4bit-V": { + "config": "1bit-K_4bit-V", + "attn_cos_sim": 0.47968902190526325, + "attn_cos_sim_std": 0.18598941169446778, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8687977194786072, + "key_bits": 144, + "val_bits": 528, + "total_bits": 672, + "total_compress": 6.095238095238095, + "n_evals": 6 + }, + "1bit-K_5bit-V": { + "config": "1bit-K_5bit-V", + "attn_cos_sim": 0.5087136526902517, + "attn_cos_sim_std": 0.2049756014656585, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9365533093611399, + "key_bits": 144, + "val_bits": 656, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "1bit-K_6bit-V": { + "config": "1bit-K_6bit-V", + "attn_cos_sim": 0.5249765266974767, + "attn_cos_sim_std": 0.19193184629905366, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9689570963382721, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "2bit-K_3bit-V": { + "config": "2bit-K_3bit-V", + "attn_cos_sim": 0.4433211386203766, + "attn_cos_sim_std": 0.19752652296671636, + "key_cos_sim": 0.9063248038291931, + "val_cos_sim": 0.8254336913426717, + "key_bits": 272, + "val_bits": 400, + "total_bits": 672, + "total_compress": 6.095238095238095, + "n_evals": 6 + }, + "2bit-K_4bit-V": { + "config": "2bit-K_4bit-V", + "attn_cos_sim": 0.4871279299259186, + "attn_cos_sim_std": 0.18471991338205915, + "key_cos_sim": 0.9063248038291931, + "val_cos_sim": 0.8687977194786072, + "key_bits": 272, + "val_bits": 528, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "2bit-K_5bit-V": { + "config": "2bit-K_5bit-V", + "attn_cos_sim": 0.5165007710456848, + "attn_cos_sim_std": 0.20428202329736883, + "key_cos_sim": 0.9063248038291931, + "val_cos_sim": 0.9365533093611399, + "key_bits": 272, + "val_bits": 656, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "1bit-K_optV-400": { + "config": "1bit-K_optV-400", + "attn_cos_sim": 0.43440664062897366, + "attn_cos_sim_std": 0.15024142809764024, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.7956738471984863, + "key_bits": 144, + "val_bits": 416, + "total_bits": 560, + "total_compress": 7.314285714285714, + "n_evals": 6 + }, + "1bit-K_optV-500": { + "config": "1bit-K_optV-500", + "attn_cos_sim": 0.4864279429117839, + "attn_cos_sim_std": 0.17132991112025142, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8773498833179474, + "key_bits": 144, + "val_bits": 516, + "total_bits": 660, + "total_compress": 6.206060606060606, + "n_evals": 6 + }, + "1bit-K_optV-600": { + "config": "1bit-K_optV-600", + "attn_cos_sim": 0.5127837210893631, + "attn_cos_sim_std": 0.18490126331804174, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9368999799092611, + "key_bits": 144, + "val_bits": 616, + "total_bits": 760, + "total_compress": 5.389473684210526, + "n_evals": 6 + }, + "1bit-K_optV-700": { + "config": "1bit-K_optV-700", + "attn_cos_sim": 0.5198984940846761, + "attn_cos_sim_std": 0.19114339427167445, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9626753330230713, + "key_bits": 144, + "val_bits": 716, + "total_bits": 860, + "total_compress": 4.7627906976744185, + "n_evals": 6 + }, + "1bit-K_optV-800": { + "config": "1bit-K_optV-800", + "attn_cos_sim": 0.5249765266974767, + "attn_cos_sim_std": 0.19193184629905366, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9689570963382721, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "2bit-K_optV-500": { + "config": "2bit-K_optV-500", + "attn_cos_sim": 0.4944220185279846, + "attn_cos_sim_std": 0.17071813167870073, + "key_cos_sim": 0.9063248038291931, + "val_cos_sim": 0.8773498833179474, + "key_bits": 272, + "val_bits": 516, + "total_bits": 788, + "total_compress": 5.197969543147208, + "n_evals": 6 + }, + "2bit-K_optV-600": { + "config": "2bit-K_optV-600", + "attn_cos_sim": 0.5209204604228338, + "attn_cos_sim_std": 0.18433921180295004, + "key_cos_sim": 0.9063248038291931, + "val_cos_sim": 0.9368999799092611, + "key_bits": 272, + "val_bits": 616, + "total_bits": 888, + "total_compress": 4.612612612612613, + "n_evals": 6 + }, + "uniform-3bit": { + "config": "uniform-3bit", + "attn_cos_sim": 0.45800502101580304, + "attn_cos_sim_std": 0.19474016746381387, + "key_cos_sim": 0.9064196844895681, + "val_cos_sim": 0.8254336913426717, + "key_bits": 400, + "val_bits": 400, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "uniform-4bit": { + "config": "uniform-4bit", + "attn_cos_sim": 0.5537208368380865, + "attn_cos_sim_std": 0.1722633167754241, + "key_cos_sim": 0.9068147341410319, + "val_cos_sim": 0.8687977194786072, + "key_bits": 528, + "val_bits": 528, + "total_bits": 1056, + "total_compress": 3.878787878787879, + "n_evals": 6 + } + }, + "benchmark_note": "1-bit keys here means per-dimension scalar quantization in the rotated spectral basis, plus a separately stored vector norm. It is not a raw sign-only key in the original basis.", + "query_capture": { + "mode": "true_model_queries", + "head_mapping": "first_query_head_per_kv_group", + "attention_mode": "causal_self_attention" + }, + "coverage": { + "preset": "small", + "n_calib": 8, + "n_eval": 4, + "seq_len": 512, + "max_cached_vectors": 256, + "layer_mode": "sampled", + "head_mode": "sampled", + "layers_evaluated": [ + 7, + 14, + 21 + ], + "heads_evaluated": [ + 0, + 1 + ] + } +} \ No newline at end of file diff --git a/results/push_095/push095_TinyLlama-1.1B-proxy-small.json b/results/push_095/push095_TinyLlama-1.1B-proxy-small.json new file mode 100644 index 0000000..31cb0e2 --- /dev/null +++ b/results/push_095/push095_TinyLlama-1.1B-proxy-small.json @@ -0,0 +1,74 @@ +{ + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "short_name": "TinyLlama-1.1B-small", + "coverage": { + "n_calib": 8, + "n_eval": 4, + "seq_len": 512, + "layer_mode": "sampled", + "head_mode": "sampled", + "max_cached_vectors": 256, + "coverage_preset": "small" + }, + "query_capture": { + "mode": "proxy_keys_as_queries", + "attention_mode": "noncausal_self_attention_proxy" + }, + "configs": { + "1bit-K_3bit-V": { + "attn_cos_sim": 0.8340571721394857, + "attn_cos_sim_std": 0.04963275925211964, + "key_cos_sim": 0.9170654565095901, + "val_cos_sim": 0.9649051427841187, + "key_bits": 80, + "val_bits": 208, + "total_bits": 288, + "total_compress": 7.111111111111111, + "n_evals": 12 + }, + "1bit-K_4bit-V": { + "attn_cos_sim": 0.8571131080389023, + "attn_cos_sim_std": 0.04418163977170474, + "key_cos_sim": 0.9170654565095901, + "val_cos_sim": 0.9902911235888799, + "key_bits": 80, + "val_bits": 272, + "total_bits": 352, + "total_compress": 5.818181818181818, + "n_evals": 12 + }, + "1bit-K_5bit-V": { + "attn_cos_sim": 0.8627720425526301, + "attn_cos_sim_std": 0.04268748699019888, + "key_cos_sim": 0.9170654565095901, + "val_cos_sim": 0.9970371077458063, + "key_bits": 80, + "val_bits": 336, + "total_bits": 416, + "total_compress": 4.923076923076923, + "n_evals": 12 + }, + "1bit-K_6bit-V": { + "attn_cos_sim": 0.8643989910682043, + "attn_cos_sim_std": 0.042387109014247304, + "key_cos_sim": 0.9170654565095901, + "val_cos_sim": 0.9988630811373392, + "key_bits": 80, + "val_bits": 400, + "total_bits": 480, + "total_compress": 4.266666666666667, + "n_evals": 12 + }, + "uniform-3bit": { + "attn_cos_sim": 0.8441474189360937, + "attn_cos_sim_std": 0.04701041290200429, + "key_cos_sim": 0.9171247283617655, + "val_cos_sim": 0.9649051427841187, + "key_bits": 208, + "val_bits": 208, + "total_bits": 416, + "total_compress": 4.923076923076923, + "n_evals": 12 + } + } +} \ No newline at end of file diff --git a/results/push_095/push095_TinyLlama-1.1B-realq-small.json b/results/push_095/push095_TinyLlama-1.1B-realq-small.json new file mode 100644 index 0000000..48bd9a8 --- /dev/null +++ b/results/push_095/push095_TinyLlama-1.1B-realq-small.json @@ -0,0 +1,75 @@ +{ + "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "short_name": "TinyLlama-1.1B-small", + "coverage": { + "n_calib": 8, + "n_eval": 4, + "seq_len": 512, + "layer_mode": "sampled", + "head_mode": "sampled", + "max_cached_vectors": 256, + "coverage_preset": "small" + }, + "query_capture": { + "mode": "true_model_queries", + "head_mapping": "first_query_head_per_kv_group", + "attention_mode": "causal_self_attention" + }, + "configs": { + "1bit-K_3bit-V": { + "attn_cos_sim": 0.47398893473049003, + "attn_cos_sim_std": 0.20592412362810367, + "key_cos_sim": 0.9170654565095901, + "val_cos_sim": 0.9649051427841187, + "key_bits": 80, + "val_bits": 208, + "total_bits": 288, + "total_compress": 7.111111111111111, + "n_evals": 12 + }, + "1bit-K_4bit-V": { + "attn_cos_sim": 0.48894164462884265, + "attn_cos_sim_std": 0.20419330170656116, + "key_cos_sim": 0.9170654565095901, + "val_cos_sim": 0.9902911235888799, + "key_bits": 80, + "val_bits": 272, + "total_bits": 352, + "total_compress": 5.818181818181818, + "n_evals": 12 + }, + "1bit-K_5bit-V": { + "attn_cos_sim": 0.49095698570211727, + "attn_cos_sim_std": 0.20493604446253103, + "key_cos_sim": 0.9170654565095901, + "val_cos_sim": 0.9970371077458063, + "key_bits": 80, + "val_bits": 336, + "total_bits": 416, + "total_compress": 4.923076923076923, + "n_evals": 12 + }, + "1bit-K_6bit-V": { + "attn_cos_sim": 0.4954307582229376, + "attn_cos_sim_std": 0.20400078514276254, + "key_cos_sim": 0.9170654565095901, + "val_cos_sim": 0.9988630811373392, + "key_bits": 80, + "val_bits": 400, + "total_bits": 480, + "total_compress": 4.266666666666667, + "n_evals": 12 + }, + "uniform-3bit": { + "attn_cos_sim": 0.5061757043004036, + "attn_cos_sim_std": 0.19893990903690034, + "key_cos_sim": 0.9171247283617655, + "val_cos_sim": 0.9649051427841187, + "key_bits": 208, + "val_bits": 208, + "total_bits": 416, + "total_compress": 4.923076923076923, + "n_evals": 12 + } + } +} \ No newline at end of file diff --git a/results/push_095/qwen1p5b_1bit_claim_check_small.json b/results/push_095/qwen1p5b_1bit_claim_check_small.json new file mode 100644 index 0000000..39c97b0 --- /dev/null +++ b/results/push_095/qwen1p5b_1bit_claim_check_small.json @@ -0,0 +1,279 @@ +{ + "model": "Qwen/Qwen2.5-1.5B-Instruct", + "coverage": { + "n_calib": 8, + "n_eval": 4, + "seq_len": 512, + "layer_mode": "sampled", + "head_mode": "sampled", + "max_cached_vectors": 256, + "coverage_preset": "small" + }, + "layers_evaluated": [ + 7, + 14, + 21 + ], + "heads_evaluated": [ + 0, + 1 + ], + "mean_d_eff_keys": 3.8392857142857144, + "mean_d_eff_values": 35.07142857142857, + "query_capture_real_q": "true_model_queries", + "query_capture_proxy": "keys_as_queries", + "attention_mode_real_q": "causal_self_attention", + "attention_mode_proxy": "noncausal_self_attention_proxy", + "configs": { + "1bit-K_3bit-V": { + "proxy_qeqk_noncausal": { + "attn_cos_sim": 0.649636780222257, + "attn_cos_sim_std": 0.2728981431678952, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8254336913426717, + "key_bits": 144, + "val_bits": 400, + "total_bits": 544, + "total_compress": 7.529411764705882, + "n_evals": 6 + }, + "real_q_causal": { + "attn_cos_sim": 0.4363829667369525, + "attn_cos_sim_std": 0.19861484805362176, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8254336913426717, + "key_bits": 144, + "val_bits": 400, + "total_bits": 544, + "total_compress": 7.529411764705882, + "n_evals": 6 + }, + "delta_attn_cos": -0.21325381348530453 + }, + "1bit-K_4bit-V": { + "proxy_qeqk_noncausal": { + "attn_cos_sim": 0.6717648692429066, + "attn_cos_sim_std": 0.2813665016012008, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8687977194786072, + "key_bits": 144, + "val_bits": 528, + "total_bits": 672, + "total_compress": 6.095238095238095, + "n_evals": 6 + }, + "real_q_causal": { + "attn_cos_sim": 0.47968902190526325, + "attn_cos_sim_std": 0.18598941169446778, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8687977194786072, + "key_bits": 144, + "val_bits": 528, + "total_bits": 672, + "total_compress": 6.095238095238095, + "n_evals": 6 + }, + "delta_attn_cos": -0.19207584733764332 + }, + "1bit-K_5bit-V": { + "proxy_qeqk_noncausal": { + "attn_cos_sim": 0.7147521277268728, + "attn_cos_sim_std": 0.2979772684157822, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9365533093611399, + "key_bits": 144, + "val_bits": 656, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "real_q_causal": { + "attn_cos_sim": 0.5087136526902517, + "attn_cos_sim_std": 0.2049756014656585, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9365533093611399, + "key_bits": 144, + "val_bits": 656, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "delta_attn_cos": -0.2060384750366211 + }, + "1bit-K_6bit-V": { + "proxy_qeqk_noncausal": { + "attn_cos_sim": 0.7450543741385142, + "attn_cos_sim_std": 0.3066504123342587, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9689570963382721, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "real_q_causal": { + "attn_cos_sim": 0.5249765266974767, + "attn_cos_sim_std": 0.19193184629905366, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9689570963382721, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "delta_attn_cos": -0.22007784744103742 + }, + "1bit-K_optV-400": { + "proxy_qeqk_noncausal": { + "attn_cos_sim": 0.5912532961616913, + "attn_cos_sim_std": 0.2508941894516606, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.7956738471984863, + "key_bits": 144, + "val_bits": 416, + "total_bits": 560, + "total_compress": 7.314285714285714, + "n_evals": 6 + }, + "real_q_causal": { + "attn_cos_sim": 0.43440664062897366, + "attn_cos_sim_std": 0.15024142809764024, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.7956738471984863, + "key_bits": 144, + "val_bits": 416, + "total_bits": 560, + "total_compress": 7.314285714285714, + "n_evals": 6 + }, + "delta_attn_cos": -0.15684665553271765 + }, + "1bit-K_optV-500": { + "proxy_qeqk_noncausal": { + "attn_cos_sim": 0.6651112387577692, + "attn_cos_sim_std": 0.2759489913732797, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8773498833179474, + "key_bits": 144, + "val_bits": 516, + "total_bits": 660, + "total_compress": 6.206060606060606, + "n_evals": 6 + }, + "real_q_causal": { + "attn_cos_sim": 0.4864279429117839, + "attn_cos_sim_std": 0.17132991112025142, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.8773498833179474, + "key_bits": 144, + "val_bits": 516, + "total_bits": 660, + "total_compress": 6.206060606060606, + "n_evals": 6 + }, + "delta_attn_cos": -0.17868329584598536 + }, + "1bit-K_optV-600": { + "proxy_qeqk_noncausal": { + "attn_cos_sim": 0.7170996777713299, + "attn_cos_sim_std": 0.2980135224683886, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9368999799092611, + "key_bits": 144, + "val_bits": 616, + "total_bits": 760, + "total_compress": 5.389473684210526, + "n_evals": 6 + }, + "real_q_causal": { + "attn_cos_sim": 0.5127837210893631, + "attn_cos_sim_std": 0.18490126331804174, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9368999799092611, + "key_bits": 144, + "val_bits": 616, + "total_bits": 760, + "total_compress": 5.389473684210526, + "n_evals": 6 + }, + "delta_attn_cos": -0.20431595668196678 + }, + "1bit-K_optV-700": { + "proxy_qeqk_noncausal": { + "attn_cos_sim": 0.7395737692713737, + "attn_cos_sim_std": 0.3046034646473958, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9626753330230713, + "key_bits": 144, + "val_bits": 716, + "total_bits": 860, + "total_compress": 4.7627906976744185, + "n_evals": 6 + }, + "real_q_causal": { + "attn_cos_sim": 0.5198984940846761, + "attn_cos_sim_std": 0.19114339427167445, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9626753330230713, + "key_bits": 144, + "val_bits": 716, + "total_bits": 860, + "total_compress": 4.7627906976744185, + "n_evals": 6 + }, + "delta_attn_cos": -0.2196752751866976 + }, + "1bit-K_optV-800": { + "proxy_qeqk_noncausal": { + "attn_cos_sim": 0.7450543741385142, + "attn_cos_sim_std": 0.3066504123342587, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9689570963382721, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "real_q_causal": { + "attn_cos_sim": 0.5249765266974767, + "attn_cos_sim_std": 0.19193184629905366, + "key_cos_sim": 0.9062935908635458, + "val_cos_sim": 0.9689570963382721, + "key_bits": 144, + "val_bits": 784, + "total_bits": 928, + "total_compress": 4.413793103448276, + "n_evals": 6 + }, + "delta_attn_cos": -0.22007784744103742 + }, + "uniform-3bit": { + "proxy_qeqk_noncausal": { + "attn_cos_sim": 0.6673308735092481, + "attn_cos_sim_std": 0.2781717269197852, + "key_cos_sim": 0.9064196844895681, + "val_cos_sim": 0.8254336913426717, + "key_bits": 400, + "val_bits": 400, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "real_q_causal": { + "attn_cos_sim": 0.45800502101580304, + "attn_cos_sim_std": 0.19474016746381387, + "key_cos_sim": 0.9064196844895681, + "val_cos_sim": 0.8254336913426717, + "key_bits": 400, + "val_bits": 400, + "total_bits": 800, + "total_compress": 5.12, + "n_evals": 6 + }, + "delta_attn_cos": -0.2093258524934451 + } + } +} \ No newline at end of file diff --git a/results/push_095/qwen1p5b_1bit_claim_check_small.md b/results/push_095/qwen1p5b_1bit_claim_check_small.md new file mode 100644 index 0000000..f217846 --- /dev/null +++ b/results/push_095/qwen1p5b_1bit_claim_check_small.md @@ -0,0 +1,83 @@ +# Qwen 1.5B 1-bit Key Claim Check + +This note compares the original favorable benchmark setup against the patched benchmark on `Qwen/Qwen2.5-1.5B-Instruct`. + +Result files: +- `results/push_095/push095_Qwen-1.5B-proxy-small.json` +- `results/push_095/push095_Qwen-1.5B-realq-small.json` + +## What Changed Versus The Initial Setup + +The original favorable setup was approximated as: +- queries replaced by keys: `Q = K` +- non-causal self-attention proxy +- attention quality measured on the proxy setup + +The patched setup used here changes that to: +- true model queries captured from the attention module +- causal attention instead of the symmetric proxy +- identical sampled coverage for both runs: + - `n_calib=8` + - `n_eval=4` + - `seq_len=512` + - `max_cached_vectors=256` + - sampled layers `[7, 14, 21]` + - sampled KV heads `[0, 1]` + +Important representation note: +- "`1-bit keys`" here still means one bit per rotated key dimension in the calibrated basis, plus a stored norm. +- It is not raw sign-only keys in the original basis. + +## Main Outcome + +The key/value reconstruction metrics barely move between the two setups, because compression is unchanged. +What moves is the attention-output cosine once the benchmark uses real queries and causal attention. + +For every tested 1-bit-key configuration, the patched benchmark is about `0.19` to `0.22` lower in attention cosine than the proxy benchmark. + +## Side-By-Side Results + +| Config | Bits | Compression | Proxy attn cos | Real-Q attn cos | Delta | +|---|---:|---:|---:|---:|---:| +| `1bit-K_3bit-V` | 544 | 7.53x | 0.6496 | 0.4364 | -0.2133 | +| `1bit-K_4bit-V` | 672 | 6.10x | 0.6718 | 0.4797 | -0.1921 | +| `1bit-K_5bit-V` | 800 | 5.12x | 0.7148 | 0.5087 | -0.2060 | +| `1bit-K_6bit-V` | 928 | 4.41x | 0.7451 | 0.5250 | -0.2201 | +| `uniform-3bit` | 800 | 5.12x | 0.6673 | 0.4580 | -0.2093 | + +## Interpretation + +- The proxy benchmark materially overstates how good the 1-bit-key setting looks. +- The drop is not caused by worse key reconstruction: + - `key_cos_sim` is the same in proxy and real-Q runs for each config. + - `val_cos_sim` is also the same. +- The difference comes from the benchmark itself, not the quantizer. + +## Extra Ablation: Real Q + Non-Causal + +An additional follow-up run was added for the missing `Real Q + non-causal` corner: +- `results/push_095/push095_Qwen-1.5B-realq-noncausal-small.json` + +For the representative `1bit-K_6bit-V` setting: + +| Setting | Attention cosine | +|---|---:| +| `Q=K + non-causal` | 0.7451 | +| `Real Q + non-causal` | 0.4124 | +| `Real Q + causal` | 0.5250 | + +What this shows on the reduced Qwen 1.5B run: +- most of the optimistic gap comes from using `Q = K` instead of true captured queries +- the causal mask is not the source of the main drop on this specific run +- the remaining missing corner is `Q=K + causal` if a full `2 x 2` attribution is needed + +## Narrower Claim Supported By This Run + +A narrower claim that is supported by this reduced Qwen 1.5B run is: + +- 1-bit rotated keys can still be competitive under a more realistic benchmark. +- But the original proxy setup overstated absolute attention quality by roughly `0.2` cosine points. +- On this run, `1bit-K_5bit-V` still beats `uniform-3bit` at the same `800`-bit budget: + - `0.5087` vs `0.4580` +- So the claim is not "1-bit keys are universally great as measured originally". +- The safer claim is "1-bit rotated keys remain viable, but the original `Q=K` proxy made them look substantially better than they do under true model queries." From 11f40da293ea4236072e5c311927444e2efc879f Mon Sep 17 00:00:00 2001 From: Your Name Date: Sun, 12 Apr 2026 20:37:54 +0200 Subject: [PATCH 2/2] add cleaned WHT vs PCA three-way experiment --- codex_wht_threeway_final_report.md | 216 ++++++ experiments/codex_wht_vs_pca.py | 651 ++++++++++++++++++ ...x_wht_vs_pca_proxy_threeway_Qwen-1.5B.json | 141 ++++ ...dex_wht_vs_pca_proxy_threeway_Qwen-1.5B.md | 56 ++ ...s_pca_realq_causal_threeway_Qwen-1.5B.json | 141 ++++ ..._vs_pca_realq_causal_threeway_Qwen-1.5B.md | 53 ++ tests/test_wht_basis.py | 62 ++ tests/test_wht_polarquant_compat.py | 67 ++ 8 files changed, 1387 insertions(+) create mode 100644 codex_wht_threeway_final_report.md create mode 100644 experiments/codex_wht_vs_pca.py create mode 100644 results/codex_wht_vs_pca/codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.json create mode 100644 results/codex_wht_vs_pca/codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.md create mode 100644 results/codex_wht_vs_pca/codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.json create mode 100644 results/codex_wht_vs_pca/codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.md create mode 100644 tests/test_wht_basis.py create mode 100644 tests/test_wht_polarquant_compat.py diff --git a/codex_wht_threeway_final_report.md b/codex_wht_threeway_final_report.md new file mode 100644 index 0000000..3a5f5f1 --- /dev/null +++ b/codex_wht_threeway_final_report.md @@ -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. diff --git a/experiments/codex_wht_vs_pca.py b/experiments/codex_wht_vs_pca.py new file mode 100644 index 0000000..6c5cdaa --- /dev/null +++ b/experiments/codex_wht_vs_pca.py @@ -0,0 +1,651 @@ +""" +Experiment 4: PCA vs WHT basis in the favorable 1-bit-key proxy setting. + +Primary setup: +- Q = K +- non-causal self-attention proxy + +Goal: +- isolate whether the strong 1-bit-key result depends on a learned per-head + PCA basis or survives under a TurboQuant-style WHT rotation. +""" + +import argparse +import gc +import json +import logging +import math +import os +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +PROJECT_ROOT = Path(__file__).parent.parent +sys.path.insert(0, str(PROJECT_ROOT / "src")) +sys.path.insert(0, str(PROJECT_ROOT / "baseline" / "turboquant_cutile")) + +from benchmark_utils import ( # noqa: E402 + causal_attention_output, + collect_eval_qkv, + noncausal_attention_output, + select_head_indices, + select_layer_indices, +) + +log = logging.getLogger("codex_wht_vs_pca") +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s - %(message)s") + +RESULTS_DIR = PROJECT_ROOT / "results" / "codex_wht_vs_pca" +RESULTS_DIR.mkdir(parents=True, exist_ok=True) +HF_TOKEN = os.environ.get("HF_TOKEN") + +MAIN_GRID = [ + ("pca", "1bit-K_6bit-V", 1, 6), + ("wht_only", "1bit-K_6bit-V", 1, 6), + ("wht_polarquant", "1bit-K_6bit-V", 1, 6), + ("pca", "1bit-K_4bit-V", 1, 4), + ("wht_only", "1bit-K_4bit-V", 1, 4), + ("wht_polarquant", "1bit-K_4bit-V", 1, 4), + ("pca", "2bit-K_6bit-V", 2, 6), + ("wht_only", "2bit-K_6bit-V", 2, 6), + ("wht_polarquant", "2bit-K_6bit-V", 2, 6), +] + + +def save_result(filename: str, data: dict) -> Path: + path = RESULTS_DIR / filename + with open(path, "w") as f: + json.dump(data, f, indent=2, default=str) + log.info("Saved: %s", path) + return path + + +def solve_lloyd_max_for_sigma(sigma, bits, max_iter=200, tol=1e-10): + if bits <= 0: + return torch.tensor([0.0], dtype=torch.float32) + if bits == 1: + c = sigma * math.sqrt(2.0 / math.pi) + return torch.tensor([-c, c], dtype=torch.float32) + from scipy import integrate + + n_levels = 1 << bits + pdf = lambda x: (1.0 / (math.sqrt(2 * math.pi) * sigma)) * math.exp(-x * x / (2 * sigma * sigma)) + lo, hi = -3.5 * sigma, 3.5 * sigma + centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)] + for _ in range(max_iter): + boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)] + edges = [lo * 3] + boundaries + [hi * 3] + new_centroids = [] + for i in range(n_levels): + a, b = edges[i], edges[i + 1] + num, _ = integrate.quad(lambda x: x * pdf(x), a, b) + den, _ = integrate.quad(pdf, a, b) + new_centroids.append(num / den if den > 1e-15 else centroids[i]) + if max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels)) < tol: + break + centroids = new_centroids + return torch.tensor(centroids, dtype=torch.float32) + + +def solve_polarquant_centroids(bits: int, d: int) -> torch.Tensor: + if bits <= 0: + return torch.tensor([0.0], dtype=torch.float32) + if bits == 1: + c = math.sqrt(2.0 / (math.pi * d)) + return torch.tensor([-c, c], dtype=torch.float32) + if bits == 2: + return torch.tensor([-1.51, -0.453, 0.453, 1.51], dtype=torch.float32) / math.sqrt(d) + return solve_lloyd_max_for_sigma(1.0 / math.sqrt(d), bits) + + +def quantize_nearest(x, centroids): + c = centroids.to(x.device) + diffs = x.unsqueeze(-1) - c + return c[diffs.abs().argmin(dim=-1).long()] + + +def load_model_tokenizer(model_name, device): + from transformers import AutoModelForCausalLM, AutoTokenizer + + needs_token = any(x in model_name for x in ["llama", "Llama", "gemma", "Gemma"]) + token = HF_TOKEN if needs_token else None + log.info("Loading %s on %s ...", model_name, device) + tokenizer = AutoTokenizer.from_pretrained(model_name, token=token) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + model_kwargs = {"token": token} + if device == "mps": + model_kwargs["torch_dtype"] = torch.float16 + model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs).to(device) + elif device == "cpu": + model_kwargs["torch_dtype"] = torch.float32 + model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs).to(device) + else: + model_kwargs["torch_dtype"] = torch.float16 + model_kwargs["device_map"] = device + model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs) + model.eval() + + cfg = model.config + n_layers = cfg.num_hidden_layers + n_kv = getattr(cfg, "num_key_value_heads", cfg.num_attention_heads) + hd = getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads) + try: + test_ids = tokenizer("test", return_tensors="pt").to(device) + with torch.no_grad(): + out = model(**test_ids, use_cache=True) + kv = out.past_key_values + try: + actual_hd = kv.key_cache[0].shape[-1] + except Exception: + actual_hd = kv[0][0].shape[-1] + if actual_hd != hd: + hd = actual_hd + except Exception: + pass + return model, tokenizer, n_layers, n_kv, hd + + +def extract_kv_layer(kv, l): + try: + return kv.key_cache[l], kv.value_cache[l] + except Exception: + pass + try: + return kv[l][0], kv[l][1] + except Exception: + pass + entry = list(kv)[l] + return entry[0], entry[1] + + +def calibrate_pca(model, tokenizer, n_calib, device, n_layers, n_kv, hd, seq_len=512): + from datasets import load_dataset + + ds = load_dataset("wikitext", "wikitext-103-raw-v1", split=f"train[:{n_calib * 5}]") + cov_keys = {(l, h): {"xtx": torch.zeros(hd, hd, dtype=torch.float64), "n": 0} for l in range(n_layers) for h in range(n_kv)} + cov_vals = {(l, h): {"xtx": torch.zeros(hd, hd, dtype=torch.float64), "n": 0} for l in range(n_layers) for h in range(n_kv)} + nd = 0 + for item in ds: + text = item.get("text", "") + if len(text.strip()) < 100: + continue + enc = tokenizer(text, return_tensors="pt", max_length=seq_len, truncation=True).to(device) + if enc["input_ids"].shape[1] < 16: + continue + with torch.no_grad(): + out = model(**enc, use_cache=True) + kv = out.past_key_values + for l in range(n_layers): + try: + k_l, v_l = extract_kv_layer(kv, l) + except Exception: + continue + k_l = k_l.float().cpu() + v_l = v_l.float().cpu() + for h in range(min(n_kv, k_l.shape[1])): + xk = k_l[0, h].double() + xv = v_l[0, h].double() + cov_keys[(l, h)]["xtx"] += xk.T @ xk + cov_keys[(l, h)]["n"] += xk.shape[0] + cov_vals[(l, h)]["xtx"] += xv.T @ xv + cov_vals[(l, h)]["n"] += xv.shape[0] + nd += 1 + if nd >= n_calib: + break + + def _eigen(cov_dict): + eigen = {} + for l in range(n_layers): + for h in range(n_kv): + n = cov_dict[(l, h)]["n"] + if n == 0: + eigen[(l, h)] = {"evec": torch.eye(hd), "ev": torch.ones(hd)} + continue + c = (cov_dict[(l, h)]["xtx"] / n).float() + ev, evec = torch.linalg.eigh(c) + eigen[(l, h)] = {"evec": evec.flip(1), "ev": ev.flip(0).clamp(min=0)} + return eigen + + return _eigen(cov_keys), _eigen(cov_vals) + + +def hadamard_matrix(n: int) -> torch.Tensor: + if n & (n - 1): + raise ValueError(f"Hadamard transform requires power-of-two dimension, got {n}") + h = torch.tensor([[1.0]], dtype=torch.float32) + while h.shape[0] < n: + h = torch.cat( + [ + torch.cat([h, h], dim=1), + torch.cat([h, -h], dim=1), + ], + dim=0, + ) + return h / math.sqrt(n) + + +def build_wht_basis(hd: int, *, seed: int) -> torch.Tensor: + base = hadamard_matrix(hd) + g = torch.Generator() + g.manual_seed(seed) + signs1 = torch.where(torch.rand(hd, generator=g) > 0.5, 1.0, -1.0).float() + signs2 = torch.where(torch.rand(hd, generator=g) > 0.5, 1.0, -1.0).float() + d1 = torch.diag(signs1) + d2 = torch.diag(signs2) + return d1 @ base @ d2 + + +def build_wht_bank(n_layers: int, n_kv: int, hd: int): + keys = {} + vals = {} + for l in range(n_layers): + for h in range(n_kv): + keys[(l, h)] = {"evec": build_wht_basis(hd, seed=17_000 + l * 257 + h), "ev": torch.ones(hd)} + vals[(l, h)] = {"evec": build_wht_basis(hd, seed=29_000 + l * 257 + h), "ev": torch.ones(hd)} + return keys, vals + + +def calibrate_wht_variances(model, tokenizer, n_calib, device, n_layers, n_kv, hd, wht_keys, wht_vals, seq_len=512): + from datasets import load_dataset + + ds = load_dataset("wikitext", "wikitext-103-raw-v1", split=f"train[:{n_calib * 5}]") + key_stats = {(l, h): {"sumsq": torch.zeros(hd, dtype=torch.float64), "n": 0} for l in range(n_layers) for h in range(n_kv)} + val_stats = {(l, h): {"sumsq": torch.zeros(hd, dtype=torch.float64), "n": 0} for l in range(n_layers) for h in range(n_kv)} + nd = 0 + for item in ds: + text = item.get("text", "") + if len(text.strip()) < 100: + continue + enc = tokenizer(text, return_tensors="pt", max_length=seq_len, truncation=True).to(device) + if enc["input_ids"].shape[1] < 16: + continue + with torch.no_grad(): + out = model(**enc, use_cache=True) + kv = out.past_key_values + for l in range(n_layers): + try: + k_l, v_l = extract_kv_layer(kv, l) + except Exception: + continue + k_l = k_l.float().cpu() + v_l = v_l.float().cpu() + for h in range(min(n_kv, k_l.shape[1])): + k_basis = wht_keys[(l, h)]["evec"] + v_basis = wht_vals[(l, h)]["evec"] + xk = k_l[0, h] + xv = v_l[0, h] + xk = xk / (torch.norm(xk, dim=-1, keepdim=True) + 1e-8) + xv = xv / (torch.norm(xv, dim=-1, keepdim=True) + 1e-8) + rot_k = xk @ k_basis + rot_v = xv @ v_basis + key_stats[(l, h)]["sumsq"] += rot_k.double().pow(2).sum(dim=0) + key_stats[(l, h)]["n"] += rot_k.shape[0] + val_stats[(l, h)]["sumsq"] += rot_v.double().pow(2).sum(dim=0) + val_stats[(l, h)]["n"] += rot_v.shape[0] + nd += 1 + if nd >= n_calib: + break + + for l in range(n_layers): + for h in range(n_kv): + kn = key_stats[(l, h)]["n"] + vn = val_stats[(l, h)]["n"] + if kn > 0: + wht_keys[(l, h)]["ev"] = (key_stats[(l, h)]["sumsq"] / kn).float().clamp(min=1e-6) + if vn > 0: + wht_vals[(l, h)]["ev"] = (val_stats[(l, h)]["sumsq"] / vn).float().clamp(min=1e-6) + return wht_keys, wht_vals + + +class BasisEngine: + def __init__(self, k_basis, k_eigs, v_basis, v_eigs, key_bits_per_dim, value_bits_per_dim, hd): + self.hd = hd + self.vk = k_basis.float() + self.vk_t = self.vk.T.contiguous() + self.vv = v_basis.float() + self.vv_t = self.vv.T.contiguous() + + k_ev = k_eigs.float().clamp(min=1e-6).cpu().numpy() + v_ev = v_eigs.float().clamp(min=1e-6).cpu().numpy() + self.k_codebooks = [solve_lloyd_max_for_sigma(float(np.sqrt(k_ev[i])), key_bits_per_dim) for i in range(hd)] + self.v_codebooks = [solve_lloyd_max_for_sigma(float(np.sqrt(v_ev[i])), value_bits_per_dim) for i in range(hd)] + self.key_total_bits = hd * key_bits_per_dim + 16 + self.val_total_bits = hd * value_bits_per_dim + 16 + + def compress_keys(self, k): + norms = torch.norm(k.float(), dim=-1, keepdim=True) + rotated = (k.float() / (norms + 1e-8)) @ self.vk + quant = rotated.clone() + for i in range(self.hd): + quant[:, i] = quantize_nearest(rotated[:, i], self.k_codebooks[i]) + return {"qr": quant, "norms": norms.squeeze(-1)} + + def compress_values(self, v): + norms = torch.norm(v.float(), dim=-1, keepdim=True) + rotated = (v.float() / (norms + 1e-8)) @ self.vv + quant = rotated.clone() + for i in range(self.hd): + quant[:, i] = quantize_nearest(rotated[:, i], self.v_codebooks[i]) + return {"qr": quant, "norms": norms.squeeze(-1)} + + def decompress(self, compressed, vt): + qr = compressed["qr"].float() + norms = compressed["norms"].float() + return (qr @ vt.to(qr.device)) * norms.unsqueeze(-1) + + def evaluate(self, q, k, v, *, attention_mode: str): + ck = self.compress_keys(k) + cv = self.compress_values(v) + k_hat = self.decompress(ck, self.vk_t) + v_hat = self.decompress(cv, self.vv_t) + attention_fn = noncausal_attention_output if attention_mode == "noncausal" else causal_attention_output + out_ref = attention_fn(q, k, v, self.hd) + out_hat = attention_fn(q, k_hat, v_hat, self.hd) + scores_ref = q.float() @ k.float().T / math.sqrt(self.hd) + scores_hat = q.float() @ k_hat.float().T / math.sqrt(self.hd) + if attention_mode == "causal" and scores_ref.shape[0] == scores_ref.shape[1]: + mask = torch.triu( + torch.full((scores_ref.shape[0], scores_ref.shape[1]), float("-inf"), device=scores_ref.device, dtype=scores_ref.dtype), + diagonal=1, + ) + scores_ref = scores_ref + mask + scores_hat = scores_hat + mask + probs_ref = F.softmax(scores_ref, dim=-1) + probs_hat = F.softmax(scores_hat, dim=-1) + eps = 1e-8 + kl = torch.sum(probs_ref.clamp_min(eps) * (torch.log(probs_ref.clamp_min(eps)) - torch.log(probs_hat.clamp_min(eps))), dim=-1).mean().item() + return { + "attn_cos_sim": F.cosine_similarity(out_ref.float(), out_hat.float(), dim=-1).mean().item(), + "attn_kl_fp16_to_compressed": kl, + "key_cos_sim": F.cosine_similarity(k.float(), k_hat.float(), dim=-1).mean().item(), + "val_cos_sim": F.cosine_similarity(v.float(), v_hat.float(), dim=-1).mean().item(), + "key_bits": self.key_total_bits, + "val_bits": self.val_total_bits, + "total_bits": self.key_total_bits + self.val_total_bits, + "total_compress": (2 * self.hd * 16) / (self.key_total_bits + self.val_total_bits), + } + + +class WHTPolarQuantEngine: + def __init__(self, k_basis, v_basis, key_bits_per_dim, value_bits_per_dim, hd): + self.hd = hd + self.vk = k_basis.float() + self.vk_t = self.vk.T.contiguous() + self.vv = v_basis.float() + self.vv_t = self.vv.T.contiguous() + self.k_codebook = solve_polarquant_centroids(key_bits_per_dim, hd) + self.v_codebook = solve_polarquant_centroids(value_bits_per_dim, hd) + self.key_total_bits = hd * key_bits_per_dim + 16 + self.val_total_bits = hd * value_bits_per_dim + 16 + + def _compress(self, x, basis, codebook): + norms = torch.norm(x.float(), dim=-1, keepdim=True) + rotated = (x.float() / (norms + 1e-8)) @ basis + quant = quantize_nearest(rotated, codebook) + return {"qr": quant, "norms": norms.squeeze(-1)} + + def compress_keys(self, k): + return self._compress(k, self.vk, self.k_codebook) + + def compress_values(self, v): + return self._compress(v, self.vv, self.v_codebook) + + def decompress(self, compressed, vt): + qr = compressed["qr"].float() + norms = compressed["norms"].float() + qr_norms = torch.norm(qr, dim=-1, keepdim=True) + qr = qr / torch.where(qr_norms > 1e-10, qr_norms, torch.ones_like(qr_norms)) + return (qr @ vt.to(qr.device)) * norms.unsqueeze(-1) + + def evaluate(self, q, k, v, *, attention_mode: str): + ck = self.compress_keys(k) + cv = self.compress_values(v) + k_hat = self.decompress(ck, self.vk_t) + v_hat = self.decompress(cv, self.vv_t) + attention_fn = noncausal_attention_output if attention_mode == "noncausal" else causal_attention_output + out_ref = attention_fn(q, k, v, self.hd) + out_hat = attention_fn(q, k_hat, v_hat, self.hd) + scores_ref = q.float() @ k.float().T / math.sqrt(self.hd) + scores_hat = q.float() @ k_hat.float().T / math.sqrt(self.hd) + if attention_mode == "causal" and scores_ref.shape[0] == scores_ref.shape[1]: + mask = torch.triu( + torch.full((scores_ref.shape[0], scores_ref.shape[1]), float("-inf"), device=scores_ref.device, dtype=scores_ref.dtype), + diagonal=1, + ) + scores_ref = scores_ref + mask + scores_hat = scores_hat + mask + probs_ref = F.softmax(scores_ref, dim=-1) + probs_hat = F.softmax(scores_hat, dim=-1) + eps = 1e-8 + kl = torch.sum(probs_ref.clamp_min(eps) * (torch.log(probs_ref.clamp_min(eps)) - torch.log(probs_hat.clamp_min(eps))), dim=-1).mean().item() + return { + "attn_cos_sim": F.cosine_similarity(out_ref.float(), out_hat.float(), dim=-1).mean().item(), + "attn_kl_fp16_to_compressed": kl, + "key_cos_sim": F.cosine_similarity(k.float(), k_hat.float(), dim=-1).mean().item(), + "val_cos_sim": F.cosine_similarity(v.float(), v_hat.float(), dim=-1).mean().item(), + "key_bits": self.key_total_bits, + "val_bits": self.val_total_bits, + "total_bits": self.key_total_bits + self.val_total_bits, + "total_compress": (2 * self.hd * 16) / (self.key_total_bits + self.val_total_bits), + } + + +def run_model( + model_name, + short_name, + device, + *, + n_calib, + n_eval, + seq_len, + layer_mode, + head_mode, + grid, + query_mode: str, + attention_mode: str, +): + model, tokenizer, n_layers, n_kv, hd = load_model_tokenizer(model_name, device) + layer_indices = select_layer_indices(n_layers, layer_mode) + head_indices = select_head_indices(n_kv, head_mode) + + pca_keys, pca_vals = calibrate_pca(model, tokenizer, n_calib, device, n_layers, n_kv, hd, seq_len=512) + wht_keys, wht_vals = build_wht_bank(n_layers, n_kv, hd) + wht_keys, wht_vals = calibrate_wht_variances( + model, + tokenizer, + n_calib, + device, + n_layers, + n_kv, + hd, + wht_keys, + wht_vals, + seq_len=512, + ) + + eval_bundle = collect_eval_qkv( + model, + tokenizer, + device, + n_eval=n_eval, + seq_len=seq_len, + head_dim=hd, + layer_indices=layer_indices, + head_indices=head_indices, + extract_kv_layer=extract_kv_layer, + split_start=n_calib * 5, + ) + + results = { + "model": model_name, + "short_name": short_name, + "device": device, + "seq_len": seq_len, + "n_calib": n_calib, + "n_eval": n_eval, + "layer_indices": layer_indices, + "head_indices": head_indices, + "primary_setup": {"query_mode": query_mode, "attention_mode": attention_mode}, + "grid": {}, + } + + for basis_type, config_name, key_bits, value_bits in grid: + metrics_list = [] + for l in layer_indices: + for h in head_indices: + kk = eval_bundle["keys"].get((l, h), []) + vv = eval_bundle["values"].get((l, h), []) + qq = eval_bundle["queries"].get((l, h), []) + if not kk or not vv: + continue + k_all = torch.cat(kk, dim=0).to(device).float() + v_all = torch.cat(vv, dim=0).to(device).float() + if query_mode == "proxy_qeqk": + q_eval = k_all + else: + if not qq: + continue + q_eval = torch.cat(qq, dim=0).to(device).float() + if basis_type == "pca": + engine = BasisEngine( + pca_keys[(l, h)]["evec"].to(device), + pca_keys[(l, h)]["ev"].to(device), + pca_vals[(l, h)]["evec"].to(device), + pca_vals[(l, h)]["ev"].to(device), + key_bits, + value_bits, + hd, + ) + elif basis_type == "wht_only": + engine = BasisEngine( + wht_keys[(l, h)]["evec"].to(device), + wht_keys[(l, h)]["ev"].to(device), + wht_vals[(l, h)]["evec"].to(device), + wht_vals[(l, h)]["ev"].to(device), + key_bits, + value_bits, + hd, + ) + else: + engine = WHTPolarQuantEngine( + wht_keys[(l, h)]["evec"].to(device), + wht_vals[(l, h)]["evec"].to(device), + key_bits, + value_bits, + hd, + ) + metrics_list.append(engine.evaluate(q_eval, k_all, v_all, attention_mode=attention_mode)) + + avg = {key: float(np.mean([m[key] for m in metrics_list])) for key in metrics_list[0]} + avg["n_evals"] = len(metrics_list) + results["grid"][f"{basis_type}:{config_name}"] = {"primary": avg} + + del model + gc.collect() + return results + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="mps" if torch.backends.mps.is_available() else "cpu") + parser.add_argument("--mode", choices=["main", "confirm"], default="main") + parser.add_argument("--model-name", default=None) + parser.add_argument("--short-name", default=None) + parser.add_argument("--n-calib", type=int, default=None) + parser.add_argument("--n-eval", type=int, default=None) + parser.add_argument( + "--grid-mode", + choices=["threeway", "legacy"], + default="threeway", + ) + parser.add_argument( + "--evaluation-mode", + choices=["proxy_qeqk_noncausal", "real_q_causal"], + default="proxy_qeqk_noncausal", + ) + args = parser.parse_args() + + t0 = time.time() + if args.evaluation_mode == "proxy_qeqk_noncausal": + query_mode = "proxy_qeqk" + attention_mode = "noncausal" + suffix = "proxy" + else: + query_mode = "real_q" + attention_mode = "causal" + suffix = "realq_causal" + + if args.grid_mode == "threeway": + suffix = f"{suffix}_threeway" + main_grid = MAIN_GRID + confirm_grid = [ + ("pca", "1bit-K_6bit-V", 1, 6), + ("wht_only", "1bit-K_6bit-V", 1, 6), + ("wht_polarquant", "1bit-K_6bit-V", 1, 6), + ] + else: + main_grid = [ + ("pca", "1bit-K_6bit-V", 1, 6), + ("wht_polarquant", "1bit-K_6bit-V", 1, 6), + ("pca", "1bit-K_4bit-V", 1, 4), + ("wht_polarquant", "1bit-K_4bit-V", 1, 4), + ("pca", "2bit-K_6bit-V", 2, 6), + ("wht_polarquant", "2bit-K_6bit-V", 2, 6), + ] + confirm_grid = [("pca", "1bit-K_6bit-V", 1, 6), ("wht_polarquant", "1bit-K_6bit-V", 1, 6)] + + if args.mode == "main": + model_name = args.model_name or "Qwen/Qwen2.5-0.5B-Instruct" + short_name = args.short_name or "Qwen-0.5B" + n_calib = args.n_calib if args.n_calib is not None else 4 + n_eval = args.n_eval if args.n_eval is not None else 1 + grid = main_grid + else: + model_name = args.model_name or "TinyLlama/TinyLlama-1.1B-Chat-v1.0" + short_name = args.short_name or "TinyLlama-1.1B" + n_calib = args.n_calib if args.n_calib is not None else 2 + n_eval = args.n_eval if args.n_eval is not None else 1 + grid = confirm_grid + + if args.mode == "main": + result = run_model( + model_name, + short_name, + args.device, + n_calib=n_calib, + n_eval=n_eval, + seq_len=512, + layer_mode="sampled", + head_mode="sampled", + grid=grid, + query_mode=query_mode, + attention_mode=attention_mode, + ) + result["elapsed_seconds"] = time.time() - t0 + save_result(f"codex_wht_vs_pca_{suffix}_{short_name}.json", result) + else: + result = run_model( + model_name, + short_name, + args.device, + n_calib=n_calib, + n_eval=n_eval, + seq_len=512, + layer_mode="sampled", + head_mode="sampled", + grid=grid, + query_mode=query_mode, + attention_mode=attention_mode, + ) + result["elapsed_seconds"] = time.time() - t0 + save_result(f"codex_wht_vs_pca_{suffix}_{short_name}.json", result) + + +if __name__ == "__main__": + main() diff --git a/results/codex_wht_vs_pca/codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.json b/results/codex_wht_vs_pca/codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.json new file mode 100644 index 0000000..0f2a115 --- /dev/null +++ b/results/codex_wht_vs_pca/codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.json @@ -0,0 +1,141 @@ +{ + "model": "Qwen/Qwen2.5-1.5B-Instruct", + "short_name": "Qwen-1.5B", + "device": "mps", + "seq_len": 512, + "n_calib": 4, + "n_eval": 1, + "layer_indices": [ + 7, + 14, + 21 + ], + "head_indices": [ + 0, + 1 + ], + "primary_setup": { + "query_mode": "proxy_qeqk", + "attention_mode": "noncausal" + }, + "grid": { + "pca:1bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.7568626503149668, + "attn_kl_fp16_to_compressed": 5.5865349769592285, + "key_cos_sim": 0.9022328853607178, + "val_cos_sim": 0.9695455233256022, + "key_bits": 144.0, + "val_bits": 784.0, + "total_bits": 928.0, + "total_compress": 4.413793103448277, + "n_evals": 6 + } + }, + "wht_only:1bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.8140885531902313, + "attn_kl_fp16_to_compressed": 1.1236141671737034, + "key_cos_sim": 0.8634132544199625, + "val_cos_sim": 0.9996187587579092, + "key_bits": 144.0, + "val_bits": 784.0, + "total_bits": 928.0, + "total_compress": 4.413793103448277, + "n_evals": 6 + } + }, + "wht_polarquant:1bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.9251146912574768, + "attn_kl_fp16_to_compressed": 0.40978433440128964, + "key_cos_sim": 0.8083983163038889, + "val_cos_sim": 0.9996177156766256, + "key_bits": 144.0, + "val_bits": 784.0, + "total_bits": 928.0, + "total_compress": 4.413793103448277, + "n_evals": 6 + } + }, + "pca:1bit-K_4bit-V": { + "primary": { + "attn_cos_sim": 0.6792900500198206, + "attn_kl_fp16_to_compressed": 5.5865349769592285, + "key_cos_sim": 0.9022328853607178, + "val_cos_sim": 0.8709203004837036, + "key_bits": 144.0, + "val_bits": 528.0, + "total_bits": 672.0, + "total_compress": 6.095238095238095, + "n_evals": 6 + } + }, + "wht_only:1bit-K_4bit-V": { + "primary": { + "attn_cos_sim": 0.8126056790351868, + "attn_kl_fp16_to_compressed": 1.1236141671737034, + "key_cos_sim": 0.8634132544199625, + "val_cos_sim": 0.9955720206101736, + "key_bits": 144.0, + "val_bits": 528.0, + "total_bits": 672.0, + "total_compress": 6.095238095238095, + "n_evals": 6 + } + }, + "wht_polarquant:1bit-K_4bit-V": { + "primary": { + "attn_cos_sim": 0.9225180447101593, + "attn_kl_fp16_to_compressed": 0.40978433440128964, + "key_cos_sim": 0.8083983163038889, + "val_cos_sim": 0.9955001274744669, + "key_bits": 144.0, + "val_bits": 528.0, + "total_bits": 672.0, + "total_compress": 6.095238095238095, + "n_evals": 6 + } + }, + "pca:2bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.762215839078029, + "attn_kl_fp16_to_compressed": 4.724521537621816, + "key_cos_sim": 0.9022488196690878, + "val_cos_sim": 0.9695455233256022, + "key_bits": 272.0, + "val_bits": 784.0, + "total_bits": 1056.0, + "total_compress": 3.878787878787879, + "n_evals": 6 + } + }, + "wht_only:2bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.9003737966219584, + "attn_kl_fp16_to_compressed": 1.2200973605116208, + "key_cos_sim": 0.9455097516377767, + "val_cos_sim": 0.9996187587579092, + "key_bits": 272.0, + "val_bits": 784.0, + "total_bits": 1056.0, + "total_compress": 3.878787878787879, + "n_evals": 6 + } + }, + "wht_polarquant:2bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.9744914670785269, + "attn_kl_fp16_to_compressed": 0.11141772475093603, + "key_cos_sim": 0.9439115822315216, + "val_cos_sim": 0.9996177156766256, + "key_bits": 272.0, + "val_bits": 784.0, + "total_bits": 1056.0, + "total_compress": 3.878787878787879, + "n_evals": 6 + } + } + }, + "elapsed_seconds": 609.498083114624 +} \ No newline at end of file diff --git a/results/codex_wht_vs_pca/codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.md b/results/codex_wht_vs_pca/codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.md new file mode 100644 index 0000000..036dc7e --- /dev/null +++ b/results/codex_wht_vs_pca/codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.md @@ -0,0 +1,56 @@ +# Three-Way Proxy Result + +Source: + +- `codex_wht_vs_pca_proxy_threeway_Qwen-1.5B.json` + +Setup: + +- model: `Qwen/Qwen2.5-1.5B-Instruct` +- `Q = K` +- non-causal attention +- sampled layers: `7, 14, 21` +- sampled heads: `0, 1` + +Variants: + +- `PCA` + Learned PCA basis plus per-dimension calibrated codebooks. +- `WHT-only` + Fixed WHT basis plus the same PCA-like per-dimension codebook structure. +- `WHT PolarQuant-style` + Fixed WHT basis plus shared `1/d`-based PolarQuant centroids and rotated-domain norm correction. + +## Results + +| 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 | + +Runtime: + +- `609.50` seconds on MPS + +## Reading + +After fixing `WHT-only` so it estimates per-coordinate WHT-space variances from normalized data, the proxy decomposition is much cleaner: + +- `PCA -> WHT-only` already improves the result substantially. +- `WHT-only -> WHT PolarQuant-style` still improves it further, but by less than before. + +So on this corrected run: + +- changing the basis plus calibrated WHT-space quantization helps a lot +- the PolarQuant-style shared `1/d` quantizer still helps on top of that + +The ranking is: + +- `WHT PolarQuant-style` > `WHT-only` > `PCA` diff --git a/results/codex_wht_vs_pca/codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.json b/results/codex_wht_vs_pca/codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.json new file mode 100644 index 0000000..879dbcc --- /dev/null +++ b/results/codex_wht_vs_pca/codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.json @@ -0,0 +1,141 @@ +{ + "model": "Qwen/Qwen2.5-1.5B-Instruct", + "short_name": "Qwen-1.5B", + "device": "mps", + "seq_len": 512, + "n_calib": 4, + "n_eval": 1, + "layer_indices": [ + 7, + 14, + 21 + ], + "head_indices": [ + 0, + 1 + ], + "primary_setup": { + "query_mode": "real_q", + "attention_mode": "causal" + }, + "grid": { + "pca:1bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.5148191054662069, + "attn_kl_fp16_to_compressed": 6.7447101672490435, + "key_cos_sim": 0.9022328853607178, + "val_cos_sim": 0.9695455233256022, + "key_bits": 144.0, + "val_bits": 784.0, + "total_bits": 928.0, + "total_compress": 4.413793103448277, + "n_evals": 6 + } + }, + "wht_only:1bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.7726523627837499, + "attn_kl_fp16_to_compressed": 0.4972422619660695, + "key_cos_sim": 0.8634132544199625, + "val_cos_sim": 0.9996187587579092, + "key_bits": 144.0, + "val_bits": 784.0, + "total_bits": 928.0, + "total_compress": 4.413793103448277, + "n_evals": 6 + } + }, + "wht_polarquant:1bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.6769297073284785, + "attn_kl_fp16_to_compressed": 0.8791843752066294, + "key_cos_sim": 0.8083983163038889, + "val_cos_sim": 0.9996177156766256, + "key_bits": 144.0, + "val_bits": 784.0, + "total_bits": 928.0, + "total_compress": 4.413793103448277, + "n_evals": 6 + } + }, + "pca:1bit-K_4bit-V": { + "primary": { + "attn_cos_sim": 0.4883749981721242, + "attn_kl_fp16_to_compressed": 6.7447101672490435, + "key_cos_sim": 0.9022328853607178, + "val_cos_sim": 0.8709203004837036, + "key_bits": 144.0, + "val_bits": 528.0, + "total_bits": 672.0, + "total_compress": 6.095238095238095, + "n_evals": 6 + } + }, + "wht_only:1bit-K_4bit-V": { + "primary": { + "attn_cos_sim": 0.7725071658690771, + "attn_kl_fp16_to_compressed": 0.4972422619660695, + "key_cos_sim": 0.8634132544199625, + "val_cos_sim": 0.9955720206101736, + "key_bits": 144.0, + "val_bits": 528.0, + "total_bits": 672.0, + "total_compress": 6.095238095238095, + "n_evals": 6 + } + }, + "wht_polarquant:1bit-K_4bit-V": { + "primary": { + "attn_cos_sim": 0.6753434439500173, + "attn_kl_fp16_to_compressed": 0.8791843752066294, + "key_cos_sim": 0.8083983163038889, + "val_cos_sim": 0.9955001274744669, + "key_bits": 144.0, + "val_bits": 528.0, + "total_bits": 672.0, + "total_compress": 6.095238095238095, + "n_evals": 6 + } + }, + "pca:2bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.5289537558952967, + "attn_kl_fp16_to_compressed": 5.93880836168925, + "key_cos_sim": 0.9022488196690878, + "val_cos_sim": 0.9695455233256022, + "key_bits": 272.0, + "val_bits": 784.0, + "total_bits": 1056.0, + "total_compress": 3.878787878787879, + "n_evals": 6 + } + }, + "wht_only:2bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.8993311822414398, + "attn_kl_fp16_to_compressed": 0.3452441009382407, + "key_cos_sim": 0.9455097516377767, + "val_cos_sim": 0.9996187587579092, + "key_bits": 272.0, + "val_bits": 784.0, + "total_bits": 1056.0, + "total_compress": 3.878787878787879, + "n_evals": 6 + } + }, + "wht_polarquant:2bit-K_6bit-V": { + "primary": { + "attn_cos_sim": 0.8964364727338155, + "attn_kl_fp16_to_compressed": 0.27388445536295575, + "key_cos_sim": 0.9439115822315216, + "val_cos_sim": 0.9996177156766256, + "key_bits": 272.0, + "val_bits": 784.0, + "total_bits": 1056.0, + "total_compress": 3.878787878787879, + "n_evals": 6 + } + } + }, + "elapsed_seconds": 506.95751333236694 +} \ No newline at end of file diff --git a/results/codex_wht_vs_pca/codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.md b/results/codex_wht_vs_pca/codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.md new file mode 100644 index 0000000..dd2dd30 --- /dev/null +++ b/results/codex_wht_vs_pca/codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.md @@ -0,0 +1,53 @@ +# Three-Way Real-Q Causal Result + +Source: + +- `codex_wht_vs_pca_realq_causal_threeway_Qwen-1.5B.json` + +Setup: + +- model: `Qwen/Qwen2.5-1.5B-Instruct` +- real captured `Q` +- causal attention +- sampled layers: `7, 14, 21` +- sampled heads: `0, 1` + +Variants: + +- `PCA` + Learned PCA basis plus per-dimension calibrated codebooks. +- `WHT-only` + Fixed WHT basis plus the same PCA-like per-dimension codebook structure. +- `WHT PolarQuant-style` + Fixed WHT basis plus shared `1/d`-based PolarQuant centroids and rotated-domain norm correction. + +## Results + +| 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 | + +Runtime: + +- `506.96` seconds on MPS + +## Reading + +After fixing `WHT-only` so it estimates per-coordinate WHT-space variances from normalized data, the realistic decomposition changed substantially: + +- `WHT-only` now beats `PCA` by a large margin. +- `WHT-only` is also slightly better than `WHT PolarQuant-style` on the 1-bit settings. +- at `2bit-K_6bit-V`, `WHT-only` and `WHT PolarQuant-style` are essentially tied, with PolarQuant slightly better on KL and `WHT-only` slightly better on attention cosine. + +So on this final realistic run, the ranking is: + +- `WHT-only` > `WHT PolarQuant-style` > `PCA` for the 1-bit points +- `WHT-only ~= WHT PolarQuant-style` > `PCA` for `2bit-K_6bit-V` diff --git a/tests/test_wht_basis.py b/tests/test_wht_basis.py new file mode 100644 index 0000000..e58143b --- /dev/null +++ b/tests/test_wht_basis.py @@ -0,0 +1,62 @@ +import math +import sys +from pathlib import Path + +import pytest +import torch +from scipy.linalg import hadamard as scipy_hadamard + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(PROJECT_ROOT / "experiments")) + +from experiments.codex_wht_vs_pca import build_wht_bank, build_wht_basis, hadamard_matrix + + +def test_hadamard_matrix_matches_scipy_reference(): + n = 8 + expected = torch.tensor(scipy_hadamard(n), dtype=torch.float32) / math.sqrt(n) + actual = hadamard_matrix(n) + assert torch.allclose(actual, expected) + + +def test_hadamard_matrix_rejects_non_power_of_two(): + with pytest.raises(ValueError, match="power-of-two"): + hadamard_matrix(12) + + +def test_wht_basis_is_deterministic_for_same_seed(): + basis_a = build_wht_basis(64, seed=1234) + basis_b = build_wht_basis(64, seed=1234) + basis_c = build_wht_basis(64, seed=1235) + + assert torch.allclose(basis_a, basis_b) + assert not torch.allclose(basis_a, basis_c) + + +def test_wht_basis_is_orthogonal_and_norm_preserving(): + basis = build_wht_basis(64, seed=7) + identity = torch.eye(64, dtype=torch.float32) + + assert torch.allclose(basis.T @ basis, identity, atol=1e-5, rtol=1e-5) + assert torch.allclose(basis @ basis.T, identity, atol=1e-5, rtol=1e-5) + + x = torch.randn(64, dtype=torch.float32) + y = x @ basis + + assert torch.allclose(torch.linalg.vector_norm(x), torch.linalg.vector_norm(y), atol=1e-5, rtol=1e-5) + assert torch.allclose(y @ basis.T, x, atol=1e-5, rtol=1e-5) + + +def test_wht_bank_returns_orthogonal_bases_for_keys_and_values(): + keys, values = build_wht_bank(n_layers=2, n_kv=2, hd=64) + identity = torch.eye(64, dtype=torch.float32) + + assert set(keys.keys()) == {(0, 0), (0, 1), (1, 0), (1, 1)} + assert set(values.keys()) == {(0, 0), (0, 1), (1, 0), (1, 1)} + + for bank in (keys, values): + for item in bank.values(): + basis = item["evec"] + assert basis.shape == (64, 64) + assert torch.allclose(basis.T @ basis, identity, atol=1e-5, rtol=1e-5) + assert torch.all(item["ev"] == 1.0) diff --git a/tests/test_wht_polarquant_compat.py b/tests/test_wht_polarquant_compat.py new file mode 100644 index 0000000..6f2c039 --- /dev/null +++ b/tests/test_wht_polarquant_compat.py @@ -0,0 +1,67 @@ +import math +import sys +from pathlib import Path + +import torch + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(PROJECT_ROOT / "experiments")) + +from experiments.codex_wht_vs_pca import WHTPolarQuantEngine, build_wht_basis, solve_polarquant_centroids + + +def _reference_polarquant_roundtrip(x: torch.Tensor, basis: torch.Tensor, bits: int) -> torch.Tensor: + d = x.shape[-1] + centroids = solve_polarquant_centroids(bits, d).to(x.device) + norms = torch.norm(x.float(), dim=-1, keepdim=True) + rotated = (x.float() / (norms + 1e-8)) @ basis + quant = centroids[(rotated.unsqueeze(-1) - centroids).abs().argmin(dim=-1)] + quant_norms = torch.norm(quant, dim=-1, keepdim=True) + quant = quant / torch.where(quant_norms > 1e-10, quant_norms, torch.ones_like(quant_norms)) + return (quant @ basis.T) * norms + + +def test_polarquant_centroids_match_tom_reference_formulas(): + one_bit = solve_polarquant_centroids(1, 128) + expected_one_bit = torch.tensor( + [-math.sqrt(2.0 / (math.pi * 128)), math.sqrt(2.0 / (math.pi * 128))], + dtype=torch.float32, + ) + assert torch.allclose(one_bit, expected_one_bit) + + two_bit = solve_polarquant_centroids(2, 128) + expected_two_bit = torch.tensor([-1.51, -0.453, 0.453, 1.51], dtype=torch.float32) / math.sqrt(128) + assert torch.allclose(two_bit, expected_two_bit) + + +def test_wht_engine_matches_tom_style_reference_roundtrip_for_keys_and_values(): + hd = 128 + k_basis = build_wht_basis(hd, seed=17) + v_basis = build_wht_basis(hd, seed=29) + engine = WHTPolarQuantEngine(k_basis, v_basis, key_bits_per_dim=1, value_bits_per_dim=4, hd=hd) + + g = torch.Generator().manual_seed(123) + keys = torch.randn(8, hd, generator=g, dtype=torch.float32) * 7.5 + values = torch.randn(8, hd, generator=g, dtype=torch.float32) * 3.0 + + key_hat = engine.decompress(engine.compress_keys(keys), engine.vk_t) + val_hat = engine.decompress(engine.compress_values(values), engine.vv_t) + + ref_key_hat = _reference_polarquant_roundtrip(keys, k_basis, bits=1) + ref_val_hat = _reference_polarquant_roundtrip(values, v_basis, bits=4) + + assert torch.allclose(key_hat, ref_key_hat, atol=1e-6, rtol=1e-6) + assert torch.allclose(val_hat, ref_val_hat, atol=1e-6, rtol=1e-6) + + +def test_wht_engine_norm_correction_preserves_zero_vectors(): + hd = 64 + basis = build_wht_basis(hd, seed=5) + engine = WHTPolarQuantEngine(basis, basis, key_bits_per_dim=3, value_bits_per_dim=3, hd=hd) + zeros = torch.zeros(4, hd, dtype=torch.float32) + + key_hat = engine.decompress(engine.compress_keys(zeros), engine.vk_t) + val_hat = engine.decompress(engine.compress_values(zeros), engine.vv_t) + + assert torch.allclose(key_hat, zeros, atol=1e-8, rtol=1e-8) + assert torch.allclose(val_hat, zeros, atol=1e-8, rtol=1e-8)