Upload README.md

README.md

@@ -0,0 +1,236 @@
+---
+license: mit
+language:
+- en
+library_name: pytorch
+pipeline_tag: text-generation
+tags:
+- text-generation
+- stream-mixer
+- linear-time
+- recurrent
+- attention-free
+- nanochat
+- small-llm
+datasets:
+- karpathy/climbmix-400b-shuffle
+- HuggingFaceTB/smol-smoltalk
+- cais/mmlu
+- allenai/ai2_arc
+- openai/gsm8k
+base_model: karpathy/nanochat
+---
+
+# Mnemo
+
+> *μνήμη — Greek for "memory"*
+
+**Mnemo** is a small attention-free language model with 117M parameters, built on the
+**Stream Mixer** architecture — a linear-time recurrent sequence mixer that uses
+multiple parallel content-routed memory streams instead of self-attention. The name
+nods to the model's recurrent memory: every layer maintains M parallel state buffers
+that "remember" content over the entire sequence without quadratic attention.
+
+The training pipeline (data, tokenizer, eval, fine-tuning) is a fork of
+[karpathy/nanochat](https://github.com/karpathy/nanochat), with the attention-based
+GPT replaced by a custom Stream Mixer block.
+
+---
+
+## Quick facts
+
+| | |
+|---|---|
+| Architecture | Stream Mixer (linear-time recurrent) |
+| Parameters | **117,179,136** |
+| Layers | 16 |
+| Hidden dim | 768 |
+| Memory streams (M) | 48 |
+| Stream state dim (D) | 96 |
+| Read heads | 6 |
+| Context length | 2048 tokens |
+| Vocab | 32,768 BPE (GPT-4-style pretokenization) |
+| Special tokens | `<\|bos\|>`, `<\|user_start\|>`, `<\|user_end\|>`, `<\|assistant_start\|>`, `<\|assistant_end\|>` |
+| Compute dtype | bf16 (Ampere+) / fp32 (T4/CPU) |
+| License | MIT |
+
+---
+
+## Architecture: Stream Mixer
+
+Mnemo's defining feature is its sequence mixer. Where a Transformer uses self-attention
+to compute pairwise interactions across tokens (cost: **O(T²)**), Mnemo uses a chunked
+parallel scan over M parallel content-routed memory streams (cost: **O(T · M · D)** —
+**linear in sequence length**).
+
+Per token *t* and per layer:
+
+1. Compute value `v[t]`, read query `q[t]`, content-router `r[t]`, and per-stream decay `α[t]`.
+2. Each memory stream `s_m` updates via `s_m[t] = α_m[t] · s_m[t-1] + r_m[t] · v[t]`.
+3. Multi-head sigmoid-gated read with QK-norm aggregates from the M streams.
+
+The full state across a layer is **(B, M, D)** — a fixed-size recurrent memory that
+the model can carry across arbitrary sequence lengths. The chunked scan implementation
+keeps numerical range bounded even for slow-decay streams.
+
+For details see the model source.
+
+---
+
+## Training
+
+### Pretraining (base model)
+
+| | |
+|---|---|
+| Corpus | [karpathy/climbmix-400b-shuffle](https://huggingface.co/datasets/karpathy/climbmix-400b-shuffle) — 88 shards |
+| Total tokens | **5.24B** (44.7× over params) |
+| Steps | 80,000 × B=32 × T=2048 |
+| Optimizer | AdamW (peak LR 1e-3, warmup 500, cosine to 1e-5, weight decay 0.1) |
+| Compute | RTX PRO 6000 Blackwell (single GPU, bf16) |
+| Wall time | **~9 hours** |
+| Best val loss | **2.9508** (perplexity ≈ 19.12) |
+
+### Supervised fine-tuning
+
+| | |
+|---|---|
+| Mixture | SmolTalk + MMLU×3 + ARC×4 + GSM8K×4 + SimpleSpelling + SpellingBee + 1000 Mnemo-branded identity convs |
+| Total conversations | ~1.09M |
+| Steps | 30,000 × B=8 × T=2048 = ~500M SFT tokens |
+| Optimizer | AdamW (peak LR 1e-4, warmup 300) |
+| Best val loss | ~1.45 (masked cross-entropy over assistant tokens only) |
+| Format | nanochat-style BOS-aligned best-fit packing with padding |
+
+### Pipeline
+
+```
+ClimbMix-400B
+   │
+   ▼
+[80k step pretrain on Stream Mixer]
+   │  best val 2.9508 @ step 79k
+   ▼
+Base checkpoint  (completes prompts)
+   │
+   ▼
+[30k step SFT on multi-task mixture]
+   │  best val ~1.45
+   ▼
+SFT checkpoint  (chat-aware — answers as Mnemo)
+```
+
+---
+
+## Capabilities and limitations
+
+### Capabilities (validated on the standard probe set)
+
+After SFT, Mnemo reliably handles:
+- Conversational dialogue in chat format (`<|user_start|>` / `<|assistant_start|>` delimiters)
+- Common factual recall: capital cities, chemical symbols, planets, basic vocabulary
+- Multiple-choice questions (MMLU/ARC format)
+- Step-by-step arithmetic in the GSM8K style (with occasional wrong answers)
+- Letter counting via manual enumeration (SpellingBee)
+- Identity Q&A consistent with its training persona
+
+### Limitations
+
+- **117M parameters** — confabulates confidently on niche facts, dates, and proper nouns
+- **No tool use, no internet, no images, no memory across sessions**
+- **2048-token context** — not pretrained for longer contexts; quality degrades past ~1500 tokens
+- **No RLHF** — outputs reflect only supervised signal; may produce inappropriate completions
+- **English only** — pretraining corpus is essentially English educational/web text
+- **Repetition prone in long generations** without `--repetition-penalty` or top-p
+- **Math beyond GSM8K-level** is unreliable; arithmetic gets shaky past 3-digit operands
+
+---
+
+## Usage
+
+### Direct loading
+
+```python
+import torch
+from tokenizers import Tokenizer
+from model import GPT
+
+tokenizer = Tokenizer.from_file('tokenizer.json')
+ckpt = torch.load('model_sft.pt', map_location='cuda')
+
+config = dict(ckpt['config'])
+config['vocab_size'] = ((tokenizer.get_vocab_size() + 63) // 64) * 64
+model = GPT.from_config(config).cuda().eval()
+
+state = {k.removeprefix('_orig_mod.'): v for k, v in ckpt['model'].items()}
+model.load_state_dict(state, strict=False)
+```
+
+### Chat CLI (recommended)
+
+```bash
+python3 chat_cli.py                   # interactive REPL
+python3 chat_cli.py -p "Who are you?"  # one-shot
+```
+
+The chat CLI handles the chat-format token wrapping (`<|bos|>` → `<|user_start|>` …)
+and stops generation cleanly on `<|assistant_end|>`. State is cached across turns
+via the recurrent state buffer — only the new tokens of each user message are
+prefilled, giving roughly **5–10× faster prefill** on multi-turn conversations than
+re-processing the entire history.
+
+### Raw inference (no chat format)
+
+```bash
+python3 infer.py -p "Photosynthesis is the process by which" --top-p 0.9 -r 1.15
+```
+
+Recommended sampling parameters (empirically tuned, see training log):
+- **Greedy / factual probes**: `-t 0`
+- **Short prose (≤500 tok)**: `-t 0.8 -k 50`
+- **Long prose (500–2000 tok)**: `-t 0.8 -k 50 --top-p 0.9 -r 1.15` (anti-loop)
+- **Diverse creative writing**: `-t 0.9 --top-p 0.85 -r 1.1`
+
+---
+
+## Probe outputs (greedy, from the base checkpoint)
+
+| Prompt | Output | Verdict |
+|---|---|---|
+| The capital of France is | "...Paris, and the capital of France is Paris" | ✓ |
+| The chemical symbol of gold is | "Au. It is a soft, silvery-white metal... good conductor of electricity and heat, useful in electrical wiring" | ✓ (Au correct; "silvery-white" wrong color) |
+| If yesterday was Friday, then tomorrow will be | "Tuesday" | ✗ (correct: Sunday) |
+| The opposite of hot is | "the cold" | ✓ |
+| The planets of the solar system are: | "Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, Neptune, Pluto" | ✓ (correct order) |
+| My favorite color is | "blue" | ✓ |
+| If 5\*x + 3 = 13, then x is | "a positive integer" loop | ✗ |
+
+**5/7 correct on the base model.** SFT improves chat-format adherence, MCQ accuracy, and persona consistency.
+
+---
+
+## Citation and acknowledgements
+
+Built on top of [karpathy/nanochat](https://github.com/karpathy/nanochat) by Andrej Karpathy.
+The Stream Mixer architecture is an attention-free experiment swapping the standard
+Transformer block for a recurrent linear-time sequence mixer.
+
+Pretraining data is [karpathy/climbmix-400b-shuffle](https://huggingface.co/datasets/karpathy/climbmix-400b-shuffle).
+SFT mixture sources: HuggingFaceTB/smol-smoltalk, cais/mmlu, allenai/ai2_arc, openai/gsm8k,
+and a custom 1000-conversation identity dataset.
+
+```bibtex
+@misc{mnemo2026,
+  title={Mnemo: A Linear-Time Recurrent Language Model},
+  author={Alvarado, Luis Miguel},
+  year={2026},
+  note={Built on karpathy/nanochat. Stream Mixer architecture.},
+  howpublished={\url{https://github.com/<your-handle>/mnemo}}
+}
+```
+
+---
+
+## License
+
+MIT. Use freely. No warranty.