baya1116/hypernet-sp-distill

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

Unofficial research portfolio. Can a small reasoning model keep doing long chain-of-thought (CoT) with a bounded KV footprint — by compressing the distant context into a handful of soft-prompt vectors instead of holding the whole transcript? This repo explores that on DeepSeek-R1-Distill-Qwen-1.5B, and ships a fast local (Apple-Silicon / MLX 4-bit) demo.

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

path	what	size
fft_out/	Final model — full fine-tuned student (student.pt) + pooler (pooler.pt)	3.55 GB + 303 MB
ce_out/	Intermediate — QLoRA r=64 adapter + pooler	295 MB + 303 MB
checkpoints/ap32_large.pt	Warm-start pooler (the original AttnPoolSP)	303 MB

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

bitsandbytes 4-bit is CUDA-only, and a 1.5B in fp16 swaps on an 8 GB Mac (2–5 s/tok). The fix is MLX 4-bit + a re-implementation of the SP-evict rollout in MLX (mlx_lm only runs the base LLM by itself). Result: **36 tok/s** on an 8 GB MacBook.

pip install mlx mlx-lm transformers torch huggingface_hub

# 1) materialise the FFT student as an HF dir, then quantise to MLX 4-bit (one-time, ~1 GB output)
python build_fft_hf.py                       # student.pt -> ./fft_hf
python -m mlx_lm convert --hf-path ./fft_hf --mlx-path ./fft_mlx4 -q --q-bits 4 --q-group-size 64

# 2a) single-turn bounded SP-evict generation
python gen_mlx.py "A farm has 30 animals (chickens and cows) with 74 legs. How many of each?" 2000

# 2b) reasoning battery (6 problems, auto-graded)
python sp_mlx.py 2000

# 2c) multi-turn 'long-CoT rally' (state carried across turns); 512 = recent raw-window size
python mt_mlx.py 512

The pooler is run in fp32 (MLX fp16 attention has a precision issue), the LLM in 4-bit. See pooler_mlx.py (a numerically-exact MLX port of the PyTorch pooler — verified to 2e-7).

Recommended inference settings (chat) — chat_mlx2.py

The base is a reasoning distill, not a chat model, so a few settings matter a lot. chat_mlx2.py bakes in the recipe (per DeepSeek-R1 guidance + this project's findings):

• 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{}."

python chat_mlx2.py 1000                              # interactive
printf 'explain a hash map\nlookup complexity?\n' | python chat_mlx2.py 1000   # scripted

Where it's good (with this recipe): explanations, math/probability, small code, and multi-turn follow-ups that reference earlier answers — e.g. it correctly does C(3,2)/C(5,2)=3/10 and, asked "more or less than 50%?", answers "30%, less than 50%" using the previous turn. Where it's weak: open-ended world knowledge (it's a 1.5B; it hallucinates facts) — pair with retrieval (runtime/). Latency is ~25–40 s/turn on an 8 GB Mac (it always thinks for a while).

Results (honest)

See RESULTS.md for the full write-up. Headlines:

• ✅ 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

The model is a 1.5B reasoner and hallucinates facts; pair it with retrieval. rag_mlx.py is a minimal demo of the runtime/ injection recipe (Context:\n{retrieved}\n\nQuestion: {q}) with a keyless Wikipedia backend, generating on the fast MLX path. The retrieved context is placed in a large raw window so facts are copied verbatim. Example (no-context → with-context): "Who wrote Neuromancer?" → hallucinates an author → "William Gibson"; "How tall is Mount Fuji?" → vague guess → "3,776 m" (copied from the snippet). Windowed-RAG (bounded long-context RAG): you can inject context longer than the raw window — the overflow is SP-compressed each step so memory stays bounded (~`rw). Measured: with a 1337-token context and rw=1000, a unique fact placed **in the last 1000 tokens** is recalled verbatim ("Dr. Elena Marsh, Reykjavik University"), but the **same fact pushed into the compressed overflow** degrades to gist only ("a paper in Computer Science" — the name is lost). So: **rank retrieval so the answer-bearing chunk sits in the last rwtokens; let supporting context spill into the gist-only compressed region.** Retrieval quality matters: if the answer sentence isn't in the snippet the grounding fails — which is exactly whyruntime/` does BGE sentence-level reranking + a 4-tier verify/refuse pipeline.

Live web search (keyless, no headless browser) — rag_web_mlx.py. runtime/web_search.py ships a DuckDuckGoSearch backend (plain urllib + bs4 against DDG's no-JS HTML endpoint) and a keyless Wikipedia fallback. rag_web_mlx.py wires DuckDuckGo → Wikipedia into the windowed-RAG generation. Findings from running it end-to-end on a laptop:

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

pip install beautifulsoup4 lxml          # no Playwright / Chromium needed
python rag_web_mlx.py "How tall is Mount Fuji?"

Conversation-memory RAG — forcibly recall a fact the SP compressed away (mem_rag_mlx.py). The SP is lossy for verbatim facts, so a specific value stated early (a code, a name, a number) is lost once it scrolls past the recent window rw into the soft-prompt. Fix: keep the full transcript in a cheap side text-memory (separate from the model's bounded KV); on each new turn, retrieve the relevant past turn(s) and re-inject them verbatim into the raw window. The model then keeps its O(1) bounded KV while gaining unbounded verbatim recall on demand. A/B (plant 3 facts → 3 filler turns → ask), rw=512, facts well past the window in the SP:

	recall the reservation code & alarm code
no memory-RAG	0/2 — hallucinates ("US123456…", "222222")
+ memory-RAG	2/2 — re-injects the planting turn → exact "QX7-2291", "4417"

python mem_rag_mlx.py 512

(Retrieval here is a simple content-word overlap over prior turns; runtime/ does it with a BGE sentence encoder. A small repetition guard + temp 0.5 tames the 1.5B's token-loop glitches.)

Tiered fallback memory — local-first, web-last (tiered_rag_mlx.py). Generalises the above into a cascading retriever, tried in order and short-circuiting at the first tier that clears a relevance bar:

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

The matched tier's text is injected verbatim into the raw window. Composite test: session 1 persists facts to disk, then a fresh session 2 (empty L1, loads L2 from disk) is probed with three questions engineered to route to L1 / L2 / L3 — routing 3/3, recall 3/3: hotel room "1408" (L1), employee ID "EMP-90832" (L2, from disk), Mount Fuji "3,776.24 m" (L3, web). Two bugs the test caught and fixed: (a) non-questions (statements, chit-chat) fell through to a web search and contaminated the reply ("1408" → "140813") → added a question-gate so only information-seeking turns retrieve; (b) the 1.5B dropped the EMP- prefix when copying → a verbatim-quote instruction + temp 0.4 keeps prefixes and punctuation intact.

python tiered_rag_mlx.py          # session1 -> disk -> fresh session2 probes L1/L2/L3

'If this were an app' — full application simulation (app_demo_mlx.py). Two sessions of a realistic on-device assistant: an onboarding session writes a profile to disk, then a fresh 'today' session auto-classifies each turn (chit-chat / fact-to-save / question), saves facts (instant "Got it — saved.", no generation), and routes questions through the tiers. Result over a mixed 8-turn conversation (greeting, fact-save, concept explanation, math, L1/L2/L3 recall):

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

Contamination guard (is_grounded + turn(retries=1)). On a retrieval-grounded turn the model is told to quote the value from the Context verbatim, so a faithful answer's salient token must appear in that Context. If no salient, novel token (capitalised / code / ≥5 chars, minus question-echoes) occurs verbatim in the retrieved context, the answer is rejected and resampled (state rolls back cleanly — conversation lives in gen/kept, the KV is rebuilt each step). It catches out-of-context ("Aki" for the CEO), empty ("(no answer)"), and question-echo ("OpenAI itself") — after retry the web turn lands on "Sam Altman". Honest limit: necessary-not-sufficient (can't catch a wrong in-context span).

Problems the simulation caught and fixed: (1) non-questions web-searched & contaminated → question-gate; (2) self-contained math got web-searched → math-gate (solve, don't search); (3) park≠parked and 1-word questions missed → fuzzy prefix overlap + trusted-local threshold; (4) fact-saves produced rambly "(no answer)" → instant ack; (5) recency bleed ("Aki" for the CEO) → "answer from the Context block only, ignore earlier conversation". See OPERATING.md for the full runbook.

python app_demo_mlx.py            # two-session app simulation + readiness report

runtime/ is a separate on-device RAG chat runtime (retrieval + refusal/verifier gating) that uses the SP model — kept here as a companion demo, not part of the core research.

archive/ holds the raw experimental scripts from development (kept for provenance).

Limitations & scope

This is an exploratory portfolio project, not a benchmarked paper. The mechanism overlaps with prior art on context compression (Gist Tokens, AutoCompressor, ICAE) and KV eviction (StreamingLLM, H2O); the contribution here is the applied study of soft-prompt compression for long-CoT reasoning continuation, the gist-vs-verbatim characterization, and the rw compression/retention trade-off, plus a working laptop-grade MLX deployment. Numbers are on small problems with small-sample (stochastic) accuracy.