HyperNet-SP — Bounded-Memory Soft-Prompt Compression for Long Chain-of-Thought

TL;DR — A tiny external "pooler" compresses everything beyond a short recent window into 32 soft-prompt vectors that the (co-trained) LLM reads as a prefix. Long single-turn CoT stays coherent and correct; multi-turn reasoning carries state across turns. The compression is lossy for verbatim facts (proper nouns, exact running numbers) — those must stay in the raw window — and there is a measurable fidelity floor vs full attention. Runs at ~36 tok/s on an 8 GB MacBook.

The idea

            distant context (everything older than the recent window)
                              │
                   ┌──────────▼───────────┐
                   │  AttnPoolSP pooler    │   32 learned query vectors cross-attend the distant
                   │  (3-layer x-attn)     │   tokens  ->  32 "soft-prompt" summary vectors
                   └──────────┬───────────┘
   query + [ 32 SP vectors ] + [ last `rw` raw tokens ] + [ current chunk ]  ──►  LLM  ──►  next token
                   ▲
   bounded buffer of distant tokens, evicted by cross-attention MASS when it exceeds `maxD`

• Soft-prompt compression — the pooler (AttnPoolSP, ~75M params) turns the distant transcript into 32 vectors. The LLM only ever attends to query + 32 SP + rw recent + current chunk, so the KV footprint is O(1) in total length, not O(length).
• Mass-based eviction — the distant buffer is capped at maxD; on overflow the tokens with the lowest pooler cross-attention mass are dropped (keep what the pooler actually uses), giving a deployment-consistent bounded memory.
• Co-training — a frozen LLM can't read the pooler's outputs well (it plateaus); the LLM is adapted to read the compressed prefix, first with QLoRA (r=64), then a full fine-tune (FFT). Objective is plain cross-entropy on DeepSeek-R1 CoT corpora (no separate teacher needed — the corpus is the distill source).

Models in this repo

Quickstart — fast local inference on Apple Silicon (MLX 4-bit)

Recommended inference settings (chat) — chat_mlx2.py

• temperature 0.6, and no system prompt — put any persona/instruction in the user turn.
• Force the reply to start with <think>\n — the distill model otherwise sometimes skips reasoning and returns an empty answer (e.g. on greetings).
• No in-loop think-budget — the model thinks to a natural </think>+EOS; bounds are an outer wall-clock timeout (240 s) + a 2000-token output cap (state hygiene, not think-shaping — it doesn't force </think>; it just stops one runaway turn from flooding the bounded window/SP and poisoning later turns), plus a degeneration guard (6 identical tokens). Empty/timed-out answers are retried (state rolled back, resample). A reasoning turn just takes its time (10–60 s, occasionally more). Note: a fully uncapped free-gen breaks multi-turn — the per-turn output cap is load-bearing.
• Large recent window rw (e.g. 1000) — bigger rw keeps more recent dialogue verbatim and markedly improves multi-turn context-following (the soft-prompt is lossy for specifics). rw is the compression↔retention dial.
• For math, append "Please reason step by step, and put your final answer within \boxed{}."

Results (honest)

• ✅ Single-turn long CoT works under compression — sets up equations, solves, verifies, boxes the answer, and (often) stops on EOS. ~80–90% correct on simple word problems (stochastic; sampling).
• ✅ Multi-turn state carry works when the running state stays in the raw window: a 5-turn dependent-arithmetic chain (8→6→12→8→4) is carried correctly with rw=512 (it breaks at rw=128, where the numbers fall into the lossy SP and are forgotten).
• ✅ Compression is genuinely active — beyond rw tokens, up to ~hundreds of distant tokens are squeezed into 32 SP each step (~18× at the tested lengths).
• ⚠️ Inherent compression floor — even a full fine-tune plateaus at deep-position CE ≈ 0.65 vs a full-attention reference of 0.56; the ~0.09 gap does not close with more LLM capacity.
• ⚠️ Lossy for verbatim facts — proper nouns / exact distant numbers die in the SP (gist, not verbatim). Keep them in the raw window. rw is the compression ↔ retention dial.
• ⚠️ Termination is stochastic — over-thinking sometimes runs long; EOS improves only with more training.

Architecture / training code

• train_sched_evict.py — AttnPoolSP, rollout_sp, mass-eviction (update_buffer/evict_topk), the building blocks.
• train_ce.py — CE / no-teacher trainer (QLoRA), bounded SP-evict, length-bucketed batching.
• train_fft.py — full fine-tune (8-bit Adam, embed+lm_head frozen), continues from the merged LoRA.
• prep_eos.py — builds the long-CoT corpus (dolphin-r1 + OpenThoughts-114k + OpenR1-Math-220k, ≥4000 tok, EOS-terminated).
• full_kv_ce.py — full-attention reference deepCE (the floor to compare against).
• pooler_mlx.py / gen_mlx.py / sp_mlx.py / mt_mlx.py / build_fft_hf.py — the local MLX demo.

External knowledge (RAG) — rag_mlx.py

• Grounding works when the answer sentence is retrieved — "How tall is Mount Fuji?" hallucinates "4,889 m" (it even claims taller than Everest) with no context, but copies the exact "3,776.24 m" straight out of the retrieved snippet.
• Use requests/urllib + bs4, not a headless browser. Bing serves headless Chromium a JS shell / dictionary-disambiguation bot-detection page (every result is a definition of the query's first word — "TALL", "WON", "Country"), so scraping it needs a fresh browser context + junk-rejecting retries and is still brittle. DuckDuckGo's html/lite endpoints server-render real organic results, so a plain HTTP GET/POST + bs4 gets clean snippets with no Chromium dependency. (PlaywrightBingSearch is still in runtime/ for when you specifically want Bing.)
• The bottleneck is retrieval, not generation. When the retrieved snippet doesn't contain the answer sentence — or the search engine ranks the wrong page — the model falls back to (sometimes wrong) parametric knowledge. This is exactly why runtime/ does BGE sentence-level reranking + a verify/refuse gate instead of feeding raw snippets. Free scrapers (DDG, Bing) also rate-limit request bursts, so both backends back off + retry.

1. L1 · same-session — turns from this conversation (in-memory),
2. L2 · prior-session — facts persisted to disk by earlier sessions (survives a fresh process),
3. L3 · web — live scrape (DuckDuckGo → Wikipedia).

• Tier routing: 5/5, every run — the memory/router architecture is deterministically solid.
• Correct answers: 4–5/5 — bounded by the 1.5B base, not the design: cross-session profile recall ("POL-55821", "Aki") and same-session recall ("B3, spot 47") are reliable; the once-flaky web turn is now stabilised by a contamination guard (below); the remaining miss is self-contained math — pure 1.5B arithmetic, no context to check against (pair with a calculator tool for exactness).

Limitations & scope