Upload README.md with huggingface_hub

README.md

@@ -0,0 +1,90 @@
+---
+language:
+- en
+license: apache-2.0
+base_model: Qwen/Qwen3-4B-Instruct-2507
+tags:
+- sparse-attention
+- ann-attention
+- distillation
+- search-projection
+- inference-optimization
+library_name: pytorch
+---
+
+# ann-sparseattention
+
+Search projections for ANN-substituted attention on
+[`Qwen/Qwen3-4B-Instruct-2507`](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507).
+
+Code: [github.com/unixsysdev/ann-sparseattention](https://github.com/unixsysdev/ann-sparseattention)
+
+## What's in this repo
+
+Per-layer linear search projections `(W_Qs, W_Ks)` of shape `[2560, 64]`,
+trained against the frozen base model's attention via contrastive +
+distillation losses. At inference these produce 64-d "search vectors" that
+let an off-the-shelf FAISS HNSW index pick the top-K keys to attend to,
+replacing dense `O(L²)` attention with `O(L·K)` ANN-substituted attention.
+
+Layers covered (pilot): `[4, 8, 12, 16, 20, 24]` — 6 of 36 layers, ~2M trainable params.
+
+## Pilot results (intermediate, step 1000 / 2000)
+
+| Step | Recall@K=128 | PPL gap (full vs ANN) |
+|---|---|---|
+| 500 | 47.4% | 1.21% |
+| 1000 | 50.7% | 0.68% |
+
+PPL gap is the primary signal — at <1% relative gap, the model's output
+quality is preserved under ANN substitution. Final-checkpoint numbers and
+the full recall@K curve over `K ∈ {64, 128, 256, 512}` will be added when
+the 2K-step pilot completes.
+
+## Files
+
+| File | What |
+|---|---|
+| `search_step_1000.pt` | Search-projection state-dict + optimizer + scheduler at step 1000 (`~11 MB`) |
+| `config.json` | Pilot hyperparams used for this checkpoint |
+
+## Loading
+
+```python
+import torch
+from transformers import AutoModelForCausalLM
+from search_module import SearchProjectionModule  # from the GitHub repo
+
+base = AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen3-4B-Instruct-2507",
+    dtype=torch.bfloat16,
+    device_map="auto",
+    attn_implementation="sdpa",
+)
+
+search = SearchProjectionModule(
+    d_model=2560, d_search=64,
+    layer_indices=[4, 8, 12, 16, 20, 24],
+    use_mlp=False,
+).to(base.device).to(torch.bfloat16)
+
+ckpt = torch.load("search_step_1000.pt", map_location="cpu")
+search.load_state_dict(ckpt["search_module"])
+```
+
+Use `inference.install_ann_attention(...)` (in the GitHub repo) to monkey-patch
+the trained layers and run with FAISS HNSW retrieval at inference time.
+
+## Training recipe
+
+- Frozen base: Qwen3-4B-Instruct-2507 (36 layers, hidden 2560, GQA 32:8).
+- Data: WikiText-103 raw, packed to 4K-token sequences.
+- 2000 steps, batch 8, lr 1e-4 (cosine, 100-step warmup), AdamW.
+- `α=β=1` (contrastive + KL distillation, both layers averaged).
+- bf16 weights, fp32 loss math.
+- SDPA attention (B200, no flash-attn package needed).
+- Liger fused RMSNorm/SwiGLU/RoPE on the frozen base.
+
+## License
+
+The search projections are released under Apache-2.0 (matching the base model).