phase-sensitive-kv-learned-500M

The reference learned compression kernel (500M-token checkpoint) from phase-sensitive-kv — an independent reproduction of Periodic Weak Spots: Phase Sensitivity from Chunked KV-Cache Compression (arXiv:2609.36322).

This is an independent reproduction, not an official release.

Model

Qwen3-0.6B backbone (28 layers, hidden 1024, 16 query heads, head dim 128) with chunked KV compression as the only architectural change: non-overlapping windows W=S=8, one KV head per layer, separate K/V branches, scalar gates (offset-specific gate maps + biases, payload maps shared across offsets), exact 16-token raw local window, one joint softmax over compressed + local memory, compressed keys RMSNorm→RoPE at (j−1)S.

Trained from scratch on FineWeb-Edu (sample-10BT), 500M tokens, seq len 2048, AdamW/cosine, 2×H100, ~2.2 h. Details and per-checkpoint evaluations: RESULTS.md.

Phase sensitivity at this checkpoint (Prefix Padding)

mean	worst phase (0–6)	gap (0–6)	boundary phase 7
0.217	0.153	13.9 pp (CI95 [12.1, 15.7])	0.12

Usage

Load with the repo code (the checkpoint is a torch dict with keys model / config / step / tokens_seen):

import torch
from pskv.model import ResCSAConfig, ResCSAForCausalLM

ck = torch.load("ckpt_499M.pt", map_location="cpu", weights_only=False)
model = ResCSAForCausalLM(ck["config"])
model.load_state_dict(ck["model"])

Evaluation: python -m pskv.eval.run_eval --ckpt ckpt_499M.pt against the frozen phase-sweep eval sets generated by the repo.

License

MIT (code); weights released under MIT as well.