supermix-expanse

Supermix Expanse

Supermix Expanse grafts, distils and fine-tunes five sources into one 132.6M-parameter PyTorch model:

Status: experimental research model. Expanse keeps and improves Archimedes' own skills, but its new code and biomedical abilities are weak after one CPU training run (see Results). "Qwen3-Coder 7B/8B" was requested; Qwen has no Qwen3-Coder at that size, so the official Qwen2.5-Coder-7B-Instruct was used. BioMedLM replaced BioMistral-7B at the owner's request.

Versions in this repo

Compared on the same held-out rows (v1's dev split; loss per reply character, comparable across tokenizers) — full tables in compare_v1_v2_v3.md:

v3 is the strongest on replay, code loss, fly and connectome, and writes code in the right one-line shape, though not yet correctly enough to pass tests. It regressed on biomedical answers (connectome rows outnumber bio rows 5:1). v2 (native latent-consolidation blocks, code in code/expanse/src/consolidation_v2.py) was slightly better than v1 only because of extra training — zeroing its new blocks removes almost none of the gain — so it is not uploaded.

Architecture

flowchart TD
    P[Prompt] --> T[Word tokenizer + embeddings<br/>10,951 words, 320-d]
    T --> L02[Trunk layers 0-2<br/>attention + MoE, 88 slots in L1-L2]
    L02 --> G{{Graft site after layer 2<br/>each graft writes through a gate that starts at 0}}
    G --> F[FlyCore<br/>22 fly brains, causal]
    G --> O[Omni v48/v38<br/>prompt encoders]
    G --> O7[Omni v7<br/>frozen 77.6M branch]
    G --> C[Male CNS core<br/>11,751 cell types, 3.83M edges]
    F & O & O7 & C --> L35[Trunk layers 3-5<br/>global attention + MoE]
    L35 --> K[Thinking core + v93 CNS] --> N[Next word]

Male CNS core. Every token's layer-2 state drives the 1,277 sensory / visual / ascending cell types; activity runs 4 recurrent steps through all 3,830,931 type→type edges of the male CNS v1.0 (122.3M typed synapses), each edge weighted by its measured input fraction and signed by its presynaptic neurotransmitter (Dale's law); the 713 descending / motor / efferent types are read out, RMS-normalised and written back through a zero-initialised gate. Wiring and signs are fixed from the data; only per-type gain, bias and leak (plus in/out projections) are learned. It runs per position, so it is causal and KV-cache exact. |W| spectral radius 0.90.

Grafted experts. New MoE slots (72 → 88 in layers 1-2) hold 96-neuron experts carved from Qwen-Coder and BioMedLM layer-0/1 MLPs: ridge maps between single-token states of teacher and student, neuron selection by domain-weighted contribution, fold-in initialisation (SwiGLU for Qwen; GELU emulated exactly as silu(1.702z)/1.702 for BioMedLM), local function-fit refinement, load-matched wake biases, and 0.1× write strength for experts whose held-out fit is weak. They were born dormant and woken over 10-50% of training.

Honest graft numbers: the teacher→student representation maps are weak (held-out R² 0.05-0.15), and only 1 of 32 grafted experts reproduces its teacher neuron group on held-out tokens (Qwen layer 0, R² 0.21); the other 31 started at 0.1× write. Most of the transferred knowledge therefore comes from distillation, not from weights.

Fixes relative to Archimedes

• FlyCore causality leak fixed. Archimedes' FlyCore.sense mean-pooled the hidden state over all positions (future tokens and padding included) and broadcast one write to every position; with the trained gate that leaks the future into every logit and makes KV-cached decoding differ from a full forward. Expanse senses per position.
• At birth every new graft is exact: with the fly graft off, Expanse reproduces Archimedes' logits with max |Δ| = 0.0 on 32 rows (receipt: supermix_expanse_grafted.receipt.json).
• The Fly Lab experience log was unavailable, so fly rows are self-distilled from the frozen original FlyCore.

Training

• Stage 1 (build): 13.6 min. Stage 2 (distil + fine-tune): 3,221 steps × batch 8 × 128 tokens, 235 min on a Snapdragon X Plus CPU (no GPU).
• Data: 13,799 train rows — Archimedes replay 3,554, self-distilled fly 1,438, Qwen code 1,421, BioMedLM bio 1,386, connectome facts 6,000.
• Losses: LM + 0.5·KL(T=2) to frozen Archimedes (cached top-32 logits) on replay + 0.2·Omni-v7 intent/domain KL + fly consensus KL and sense MSE.
• LR 3e-5 trunk / 2e-4 grafts (new embedding rows and grafted router rows boosted to the graft LR), cosine, 5% warm-up, clip 1.0.
• Teacher data quality controls: code answers kept only if they pass our own tests in an AST-whitelisted, isolated subprocess sandbox (92.8% pass rate); BioMedLM definitions kept only if Qwen-Coder judged them accurate (94% accepted — a lenient judge); PubMedQA answers kept only when they matched the expert label, then capped because BioMedLM answers almost everything "yes" (on held-out PubMedQA it scores 59.3% vs a 60.5% always-yes baseline).

Results

Dev loss per source (nats/token on reply tokens, held-out dev rows; connectome dev = held-out cell types):

Agreement with Omni v7: intent 72.7%, domain 73.1% (from 4% / 10%). Final gate magnitudes: CNS core 0.086, Omni v7 0.076.

Sample answers (greedy):

• whats the impulse from 25 N acting for 4 s → impulse = force x time, 20 x 4 = 80, 5 x 4 = 20, the impulse is 100 newton seconds, total 100 ✅
• Python: r = sum(range(1, 6)) What is r? → range counts from 1 up to 5, … r is 15, total 15 ✅
• Which neurotransmitter does the cell type LC4 use in the male CNS? → answers a different connectome question (synapse count) ❌
• What is a hepatocyte? / Write a Python function that returns the sum of a list. → not usable ❌

Held-out generation (greedy, 150 items per metric; full report eval_report.md):

Ablations (Expanse dev loss with one component removed at inference):

The male-CNS core is the component the model relies on most, on every source including code, and a matched random graph loses 65-90% of its contribution — the trained model depends on the fly's specific wiring. Omni v7 helps slightly; the grafted donor experts and FlyCore contribute nothing measurable (the new knowledge came from distillation).

What is and is not shown. A gate opening shows training uses the connectome core; it does not show that the fly's wiring matters. The proper test is a matched training run on the degree/sign-preserving rewired graph (code included: connectome_full.rewired_graph) and a no-core control.

Usage

git clone https://huggingface.co/Kai9987kai/supermix-expanse
cd supermix-expanse/code/expanse
pip install torch sympy gradio
python expanse_chat.py --checkpoint ../../supermix_expanse_v3.pt     # local chat UI (v3)

import sys, torch; sys.path.insert(0, "src")          # from code/expanse
import expanse_core as ec
model, tok, payload = ec.load_expanse("../../supermix_expanse_v3.pt")  # v3 (BPE) or v1 (word)
q = "whats the impulse from 40 N acting for 6 s"
ids, _ = tok.encode_turn(q, None)
out = ec.greedy_decode(model, torch.tensor([ids]), max_new_tokens=80,
                       omni_features=ec.arch.OmniCore.featurize([q]),
                       omni7_state=model.omni7_state_for([q]))
print(tok.decode(out))

The checkpoint is a pickled PyTorch payload (torch.load(weights_only=False)), not a Transformers model; load only if you trust it. Single-turn, 128-token context; v3 uses a byte-level BPE tokenizer, v1 a word-level vocabulary.

Files

• supermix_expanse_v3.pt — v3, recommended (schema supermix-expanse-v1 with a BPE tokenizer); supermix_expanse.pt — v1. Both self-contained (frozen Omni v7 net and connectome CSR buffers included)
• compare_v1_v2.md, compare_v1_v2_v3.md (+ .json) — same-row comparisons; supermix_expanse_v3*.receipt.json — v3 retokenisation and training receipts
• supermix_expanse.receipt.json, supermix_expanse_grafted.receipt.json — training and build receipts (per-expert graft fits, vocabulary lift R², function-preservation check, full training history)
• code/expanse/ — build, train, eval, teacher-data, chat and connectome code (DESIGN.md is the spec); code/supermix-archimedes/ and code/external/omni_v7/ — the source it depends on
• data/ — the teacher corpora actually used (code, cleaned bio, connectome facts), data/v3/ (fresh verified problems + connectome v3 rows), the unbiased PubMedQA eval set, and data receipts

License and attribution

Composite — see LICENSE.md. In short: code MIT (Kai Piper); the weights contain parameters derived from BioMedLM and were trained on its outputs, so the BigScience OpenRAIL-M use restrictions apply to this model and its derivatives — in particular it must not be used to provide medical advice or medical results interpretation; Qwen2.5-Coder-7B-Instruct is Apache-2.0; the male CNS connectome data (Janelia FlyEM and collaborators) is CC BY 4.0; PubMedQA is MIT.

Citations

• Male CNS connectome v1.0 — Janelia FlyEM Project Team and the Cambridge Drosophila Connectomics Group, https://male-cns.janelia.org/; natverse malecns R package.
• Bolton et al., BioMedLM: A 2.7B Parameter Language Model Trained On Biomedical Text (2024).
• Hui et al., Qwen2.5-Coder Technical Report (2024).
• Jin et al., PubMedQA: A Dataset for Biomedical Research Question Answering (EMNLP 2019).
• Lappalainen et al., Connectome-constrained networks predict neural activity across the fly visual system (Nature 2024) — the fixed-wiring / learned-parameters recipe.