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