MaleCNS-SpikeWhale-Anatomy-LM

MaleCNS x SpikeWhale anatomy LM (step 2855)

Research artifact. Undertrained. Not a usable language model.

Hard-LIF language SNN from scratch. Recurrent matrix is the Janelia / Google MaleCNS v1.0 fruit-fly CNS, not a dense 4x1488 SpikeWhale stack. LIF cell is the SpikeWhale primitive (lif.py: hard spikes, fast-sigmoid slope 20, soft reset, DERF clamp 30). Synapse values train with nnz-only autograd through the sparse CSR.

Latest language checkpoint: real_seq2048_b4 step 2855 (checkpoints/real_seq2048_b4_step2855.pt). Seq-2048 / batch-4 continuation of step 2500. This leg: 350 steps x 8192 tokens/step ~ 2.87M tokens. Total ~5.4M. Greedy still latches. The the the the stutter at 2500 became a longer repeated clause ("It is after the ...). Different failure, not a fix. This is a measurement, not SpikeWhale-216M quality.

Older checkpoints stay: checkpoints/step2500.pt (seq 256; generations and curve below) and checkpoints/step1000.pt (W-ablation snapshot).

Connectome export / Dale numbers below are from step 2948.

Connectome (this graph)

Wiring is MaleCNS v1.0 from FlyEM (HHMI Janelia), the University of Cambridge, the MRC LMB, and Google Research. Hugging Face Papers only auto-tags arXiv IDs. This paper is in Cell, so the Hub will not grow an arxiv: chip.

Connectome data: CC-BY 4.0.

Architecture (this checkpoint)

token id
  -> embed (64)
  -> current on: 5 sensory-group winners + top-32 of 256 highest-out-degree cells (drive 0.8)
  -> 4 LIF ticks, hold current, rec = (W @ spikes) / sqrt(in_degree) * rec_scale(5)
  -> Kenyon spikes re-enter W (mask_kenyon=false); spike-mass cap 80
  -> tanh(pre-reset)[Kenyon U hubs] -> LayerNorm -> Linear 128 -> vocab 16512

• kc_drive = 0. NT gains frozen. Trainable: sparse W, embed, hub_in, head.
• Optimizer: MuonEq-R on 2D except out; AdamW on 1D / embed / out. Peak Adam 1e-3, Muon 5e-3.
• Data: 40% FineWeb-Edu / 40% DCLM / 20% IFM (pretraining_config.json; 9 streams - the IFM share is 7 behavior configs at 0.0286 each).
• Tokenizer: SpikeWhale byte-level length-max, vocab 16512.
• code/sparse_mm_fire.py: fire-only forward (scatter firing cells); d_spikes = W^T @ gout via cuSPARSE (exact LIF). Config checkpoint_every: 0 = no activation checkpoint (full tape).

Training status

Step 2855 (seq-2048 continuation) - latest language ckpt

Nothing latched and nothing broke. 350 steps at 8x the tokens per step did not move CE out of the band that 1500-2500 already sat in.

Step 2500 (seq 256)

Kenyon/readout did not latch. Language did not leave the weak-unigram / broken-English band. 1.5k extra steps after 1000 moved CE 4.89 -> 4.78. That is not a ceiling proof; it is also not a generation win.

What actually trains

Gradient reaches a synapse only if its presynaptic cell fired (sparse_mm_fire.py computes grad_values for active presynaptic cells only). Measured against the conversion weights at step 2948:

~78% of the "trainable" synapses have never received a gradient. The edges that did move changed by more than their own original magnitude.

Instantaneous participation is far lower than cumulative. Counting spikes per neuron over a 128-token window (logs/participation_2500_vs_2855.json):

~1.5% of the connectome is active at any moment, and the spike mass concentrates further into the readout population as training goes on. The effective model is a ~2.5k-neuron subnetwork inside a 167,565-neuron graph.

This is not a reachability limit. From the sensory + hub injection sites, one synaptic hop already reaches 159,113 neurons (95%), and 165,378 (98.7%) are both drivable and able to influence the readout within the 4 ticks. The graph is small-world; the sparsity is dynamics, not wiring.

Spike-mass cap: what pins participation

GradThroughSNN.step rescales every spike by min(1, max_spikes / mass). With ~268 cells firing and max_spikes: 80, the scale is ~0.30 and total spike mass is pinned: recruitment is zero-sum, because a newly firing cell dilutes the ones already carrying signal. Ramping the cap (logs/spike_cap_experiment.json):

The cap was the binding constraint. 80 -> 400 tripled the number of firing neurons, raised whole-network activity 3.6x and doubled the gradient reaching W, at stable logit_rms. 2000 broke training outright: grad_norm 2.4 -> 48.3, CE above uniform. Note 0.461 is close to the 0.45 where failure mode 2 lives.

What it did not buy: CE. Over 150 steps at cap 400 the mean went 5.017 -> 4.939 (best single step 4.468) - back into the pre-cap plateau band, not below it. More of the network participates and more synapses receive gradient; that has not yet become better prediction.

Dale's law and the exported connectome

Training is free to flip a synapse's sign, which makes the graph no longer neurotransmitter-consistent. By step 2948, 1,090,918 synapses had flipped sign (19.3% of the edges that changed, 4.3% of all edges).

export_trained_brain.py recombines a trained checkpoint's synapse vector with the frozen artifact's wiring and writes a plain connectome SNN - same neurons, same edges, same identity matrix, trained synapses, no language-model parts - loadable by ConnectomeSNN.load():

python export_trained_brain.py \
  --ckpt <trained_checkpoint>.pt \
  --out male_cns_spikewhale_trained.pt          # add --dale for the constrained one

--dale clamps every synapse back onto the sign its neurotransmitter gave it (violating ones go to 0), so the exported brain stays Dale-compliant. The trainer also takes dale_constraint: true, which re-projects after every optimizer step - that has not been run yet.

The constrained brain fires more, because zeroing sign-violating synapses removes inhibition that training created by flipping excitatory synapses negative.

Both exports still drive language when paired with the step-2948 head. Greedy, 32 new tokens, same prompts, same decode:

                  AS-TRAINED (0 zeroed)                DALE-PROJECTED (1,090,918 zeroed)
The               The to the |-------------            The therect then them to bet is again,
                                                       them therect then, them to there is after

Once upon a time  Once upon a times.                   Once upon a times, therecting then then
                  (blank lines to the limit)           then thenks.

The fruit fly     The fruit fly brain.                 The fruit fly brain, then then them to
brain             (blank lines to the limit)           because the |------------------

Greedy on the as-trained brain collapses into whitespace or dashes within a few tokens on all three prompts; the Dale-projected one keeps emitting word-shaped tokens for the full 32. Neither is English.

That is not the constrained brain modelling language better. On identical batches (logs/dale_vs_astrained_eval.json):

Projection costs 0.22 nats of prediction. What improves is greedy degeneracy only - 1.8x more distinct tokens, far less filler - which is a decoding artifact, not accuracy. Under real language input the constrained brain fires less, not more, so the pulse test above does not explain the generations: an earlier "more excitable" reading of this was wrong.

The likelier mechanism: the head spent 2,948 steps fitting one feature distribution and learned a strong period-then-whitespace attractor; any large perturbation of W knocks the features off that memorised point, breaking the attractor while also hurting accuracy. The control that would settle it - zeroing 1,090,918 randomly chosen synapses of matched magnitude - has not been run.

Sampled (T=0.8, top-k 50) from the Dale-projected brain:

Once upon a time thems are really, and the reven anly to a back that the said, deal, a cleanks i
The fruit fly brain is then a lonk-sing but themss to bet it was bet is a site wasn in if you

The head was trained against unconstrained weights, so all of this measures how a head copes with a brain that shifted under it - not the quality of a model trained with the constraint on. The trainer supports that (dale_constraint); it has not been run. Full text: logs/generate_connectome_step2948.log.

What survives as a result: removing every Dale-violating learned update costs only 0.22 nats. The sign-illegal part of what training learned is real, but small.

W ablation (step 1000 only)

Same 4 batches. Not a trained scramble control. Not re-run at 2500. File logs/eval_w_ablation_step2855.json is in the repo.

Trained synapses are used at eval.

Generations (still not good)

Step 2855

Greedy no longer stutters single words. It repeats a clause instead. Two of the five prompts collapse straight into dash / asterisk runs. Sampled text is word-shaped fragments, not English.

PROMPT: The
GREEDY: The, then the "It is the "It is again, the "It is after the "It is after the
SAMPLE: The) and a looking, and graphructs
Socing. Opt-Put, and we'p the entibolor, somethip

PROMPT: Once upon a time
GREEDY: Once upon a times, there are the "It is after the "It is aftly the -derstand and there is a s

PROMPT: In 2020, scientists discovered
GREEDY: In 2020, scientists discovered, the |------------------------------

Against step 2500: the the the and there is there is are gone, and are now agrees after a plural - a longer repeated unit, not a working sentence. Full log: logs/generate_step2855.log.

Step 2500

Greedy still latches. Sampled is slightly more word-like than 1k, still not English.

PROMPT: Once upon a time
GREEDY: Once upon a times, there is the | |---...
SAMPLE: Once upon a time about Artue. It work. ...

PROMPT: The fruit fly brain
GREEDY: The fruit fly brain, there is there is the |---...

Full log: logs/generate_step2500.log. Compare logs/generate_step1000.log.

What this does show

It can learn something. CE 9.71 -> ~4.8 by 2.56M tokens. w_grad live. Rate in band. Step-1000 ablation: destroying W wrecks CE.

Raising the spike cap engages more of the graph but has not moved CE. Cap 400 triples the firing population and doubles w_grad; 150 steps there ended at CE 4.94, inside the same band.

Then it stops. From step ~1500 through 2855 - across an 8x sequence-length change, an 8x tokens-per-step change, and a switch to padding-free packing - CE sits at 4.8-4.9. A fit over 1500-2500 gives -0.0009 CE per 100 steps: flat. Whatever caps this is not step count and not batch size.

That learning is weak language. Greedy at 2500 is still degenerate. Last-50 mean CE ~4.84. Best single batch 3.96 is noise, not a 4-break (need a 100-step mean). 2.56M tokens is a probe, not a pretrain.

• Shows: sparse MaleCNS + SpikeWhale LIF accepts gradient and beats chance.
• Does not show: more steps => fluent text, or a trained scramble control, or fly-as-attention helping a transformer.

Known failure modes (already hit)

1. 1 tick/token: sensory atlas does not reach Kenyon. Anatomy is multi-hop.
2. Seq 512 + strong rec: Kenyon clique latches.
3. 1.1 current on Kenyon: greedy the the the.
4. Muon 0.02 on out: logit blowup.
5. Fly graph as attention (prefix-sum + hub CSR instead of MHA): ~50M Byrne-looped GPT, ~120M tokens, CE stuck ~6.4. Not this checkpoint.
6. Triton per-neuron d_spikes: ~110 s/step vs ~75 s cuSPARSE. Reverted; fire-only is forward only.
7. Long windows with one document per window: at seq 2048+ the single-document blocks() pads most windows, and pad tokens cost full LIF compute before being masked out. Pack documents instead.
8. checkpoint_every: 0 at long seq: the full tape is fine at seq 256, but at seq 8192 the per-token boundary state alone is ~44 GB. Segment it (checkpoint_every: 8).

Files

Load / generate

PYTHONPATH=code python generate.py \
  --ckpt checkpoints/real_seq2048_b4_step2855.pt \
  --device cuda --n-new 64 --temperature 0.8 --top-k 50

generate.py expects male_cns_spikewhale.pt at ../../output/male_cns_spikewhale.pt relative to experiments/anatomy_lm. In this snapshot:

mkdir -p output experiments/anatomy_lm/output/real_seq256_hubs atlas
cp checkpoints/male_cns_spikewhale.pt output/
cp checkpoints/step2500.pt experiments/anatomy_lm/output/real_seq256_hubs/latest.pt
cp -r atlas experiments/anatomy_lm/
cp config.json experiments/anatomy_lm/

Resume training

Same run name loads latest.pt. Optimizer moments are not in the ckpt.

cp checkpoints/real_seq2048_b4_step2855.pt \
   experiments/anatomy_lm/output/real_seq2048_b4/latest.pt
python -u experiments/anatomy_lm/train_anatomy.py \
  --start-training --device cuda --condition real --run-name real_seq2048_b4
# expect {"resume_from": 2855, ...}

The seq-2048 leg ran with sequence_length: 2048, batch_size: 4, gradient_accumulation: 1, checkpoint_every: 8 and contiguous document packing. At seq 256 the old config.json still applies (checkpoint_every: 0, run name real_seq256_hubs, resume from step2500.pt).

Do not copy an older checkpoint over latest.pt unless you intend to rewind.

License

Code: Apache-2.0, consistent with SpikeWhale-SNN-216M.

MaleCNS data: CC-BY 4.0 (Janelia / Cambridge / MRC LMB / Google Research). Cite Berg et al. 2026 if you use the graph.

@article{berg2026malecns,
  title   = {Sexual dimorphism in the complete {Drosophila} male central nervous system connectome},
  author  = {Berg, Stuart and Beckett, Isabella R. and Costa, Marta and Schlegel, Philipp and Januszewski, Micha{\l} and Marin, Elizabeth C. and Nern, Aljoscha and others},
  journal = {Cell},
  volume  = {189},
  number  = {18},
  pages   = {5504--5526.e15},
  year    = {2026},
  doi     = {10.1016/j.cell.2026.08.015}
}