| --- |
| 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. |
|
|
