Add README.md

README.md

@@ -0,0 +1,78 @@
+---
+language:
+- en
+library_name: pytorch
+pipeline_tag: text-generation
+tags:
+- speculative-decoding
+- memory-layers
+- n-gram-memory
+- lngram
+- interpretability
+- research
+datasets:
+- OptimalScale/ClimbMix
+---
+
+# Self-Drafting Latent Memory: model collection
+
+Weights for the paper *"Self-Drafting Latent Memory: What an n-gram Memory Learns, and When It Can Draft for Its Own Model"*.
+
+These are **research checkpoints**: 370M-parameter decoder language models with a route n-gram memory ([Lngram](https://github.com/zyaaa-ux/Lngram) v2) at layer 1. Each was trained on 1B tokens of ClimbMix (2B for the dense control). They are not instruction-tuned and are not intended for deployment.
+
+In the self-drafting variants, the memory's router reads only token embeddings. A small *draft head* can therefore read the memory table **without running the transformer backbone**, and propose tokens for speculative decoding of its own model. The draft head is trained jointly with the language model.
+
+## Models
+
+| Folder | Model | Description | Non-embedding params | Val loss (1B tok) |
+|---|---|---|---|---|
+| [`sb-run1/`](sb-run1) | SB, run 1 | Baseline: Lngram v2 at layer 1, hidden-state router, 184 routes | 369.78M | 2.7127 |
+| [`sb-run2/`](sb-run2) | SB, run 2 | Same recipe, second training run (replicate) | 369.78M | 2.7107 |
+| [`sd-b/`](sd-b) | SD-B | Token-embedding router + jointly trained draft head, 182 routes | 369.93M | 2.7159 |
+| [`sd-c/`](sd-c) | SD-C | SD-B + 256-channel diagonal SSM feeding router and draft query, 181 routes | 369.88M | 2.7111 |
+| [`d270/`](d270) | D270 | Dense control: no memory, same active params per token, **2B tokens** | 269.80M | 2.5147 (2B) |
+| [`posthoc-drafters/`](posthoc-drafters) | post-hoc drafters | Token router + draft readout fitted to frozen SB run 1, without / with SSM | ~2–2.5M each | — |
+
+## Key results
+
+- **Self-drafting is free.** Loss changes by +0.0004 to +0.003 nats against the baseline, and the two baseline runs themselves differ by 0.002. There's no significant difference on HellaSwag, MMLU-cloze, SciQ, WinoGrande, PIQA or ARC-Easy.
+- **Joint training makes a far better drafter than post-hoc fitting.** Top-1 agreement with the target is 0.45 (SD-C) vs 0.30, and accepted length roughly doubles.
+- **Acceptance at γ=4 with code-keyed online counts:** 2.5–2.6 tokens per verification step. That beats prompt lookup on open-ended text. A hybrid (prompt lookup, else the draft head) beats prompt lookup on every model and text type.
+- **The drafter's own confidence predicts agreement.** Gating at 0.5 gives 88% agreement on 18% of positions.
+
+All numbers, logs and analysis code are in the GitHub repository [MorbiusPlexus/lngram-decode](https://github.com/MorbiusPlexus/lngram-decode). Training curves for every model are in the W&B project [morbius-wave3](https://forge.coreweave.com/wandb/kishor-iiit-hyderabad/morbius-wave3).
+
+## Usage
+
+The checkpoints use the `ffnmem` code from [MorbiusPlexus/lngram-decode](https://github.com/MorbiusPlexus/lngram-decode): a custom architecture, so it isn't loadable through `transformers`.
+
+```python
+import torch
+from ffnmem.infer import load_model, greedy_generate
+
+model, cfg = load_model("sd-c/model.pt", device="cuda")   # weights + resolved config
+prompt = torch.tensor([[...token ids...]], device="cuda")  # tokenizer: data/tokenizer/tokenizer.json in the repo
+tokens, _ = greedy_generate(model, prompt, n_new=64, amp_dtype=torch.float16)
+
+# Self-drafting models: draft logits from token ids alone (no backbone forward)
+b = torch.zeros(prompt.shape[1], dtype=torch.long, device="cuda")
+t = torch.arange(prompt.shape[1], device="cuda")
+draft_logits = model.draft_logits(prompt, b, t)            # SD-B / SD-C only
+```
+
+Each `model.pt` is a dict `{"model": state_dict, "config": resolved config, "step": int, "params": counts}`. `config.json` duplicates the config for inspection.
+
+## Training data and procedure
+
+- **Data:** ClimbMix, tokenised with a 32,768-token byte-level BPE (in the repository).
+- **Batch and schedule:** global batch 256 × 2048 tokens; 1,907 steps (1B tokens); cosine learning rate 6e-4 → 6e-5 with 100 warmup steps.
+- **Optimisation:** AdamW (0.9, 0.95), weight decay 0.1; memory tables use a 4× learning rate without decay.
+- **Precision:** fp16 with GradScaler.
+- **Draft loss:** λ = 0.1 × KL(sg(p) ‖ q) on 25% of positions.
+
+## Limitations
+
+- **Small models:** 370M parameters trained on 1B tokens, with performance near chance on MMLU and WinoGrande.
+- **No safety work:** not aligned or safety-tuned; outputs can be incorrect or inappropriate.
+- **Seeds:** one seed per model, except the baseline, which was trained twice.
+- **Speed:** measured offline under greedy decoding. The released live speculative loop isn't optimised and doesn't yet realise a speedup.