afoolishman/self-drafting-latent-memory

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 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, run 1	Baseline: Lngram v2 at layer 1, hidden-state router, 184 routes	369.78M	2.7127
sb-run2/	SB, run 2	Same recipe, second training run (replicate)	369.78M	2.7107
sd-b/	SD-B	Token-embedding router + jointly trained draft head, 182 routes	369.93M	2.7159
sd-c/	SD-C	SD-B + 256-channel diagonal SSM feeding router and draft query, 181 routes	369.88M	2.7111
d270/	D270	Dense control: no memory, same active params per token, 2B tokens	269.80M	2.5147 (2B)
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. Training curves for every model are in the W&B project morbius-wave3.

Usage

The checkpoints use the ffnmem code from MorbiusPlexus/lngram-decode: a custom architecture, so it isn't loadable through transformers.

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.