fernandofernandes/fly-wordbrain-rank64

Fly Wordbrain — rank 64

A 3,956,469-parameter language model whose entire recurrent layer is the measured wiring of a fruit fly brain. On a held-out population it scores 1.039 nats better in cross-entropy and 6.05 accuracy points higher than the 52,756,661-parameter model it was derived from — while being 13.3× smaller.

Two sets of weights live here. corpus-10k/ is trained on 10,000 stories and is the one you want (audit CE 2.9493, 37.43% top-1). The root files are the original 1,000-story weights (CE 3.2802), kept because every earlier experiment compares against them. Same architecture, same frozen connectome; only the training corpus differs.

Hear it read aloud: Fly Recital · Research and receipts: fly-wordbrain

What it is

A reproduction of the ngxson Fly LLM architecture with its output readout replaced by a low-rank factorization. One recurrent step per token over all 49,393 neurons and 9,050,172 measured synapses:

drive  = concat over 8 delay slots j of  wte[token(t-j)] @ in_proj[j]
x      = 0.1·x + 0.9·tanh( gain · (rec_gain · (W·x) + drive) + bias )
logits = layer_norm(x) @ head_a.T @ head_b.T          # 49,393 → 64 → 1,024

W is the connectome, frozen. No synaptic weight was modified — the frozen-buffer digests in manifest.json are byte-identical to the reference model's, across all seven graph buffers. No edge was changed and none was rewired.

The neurons are not inert, though. Each of the 49,393 carries a learned input gain, recurrent gain and bias — 148,179 parameters, part of the reference architecture rather than an addition here — and training moves them a long way. Against the seed-rebuilt initialization: gain 58.66% relative L2, rec_gain 12.51%, and the effective per-neuron scaling gain × rec_gain 50.61%. Because rec_gain multiplies a neuron's entire incoming sum, it rescales that neuron's synapses by one factor and cannot alter their relative strengths or signs. Structure preserved, per-neuron scale learned.

Component	Parameters
Encoder (embedding + 8 input projections)	482,816
Neuron gain, recurrent gain, bias	148,179
Output LayerNorm	98,786
Factorized readout (rank 64)	3,226,688
Total	3,956,469

For comparison, the reference model spends 50,578,432 parameters — 95.9% of everything it learns — on a single full-rank output matrix.

Held-out scores

200 fresh TinyStories, 45,059 next-token targets, a population that selected no checkpoint. Paired whole-story bootstrap, 10,000 resamples.

Model	Parameters	CE	PPL	Top-1	ΔCE vs reference (95% CI)
ngxson reference	52,756,661	3.9882	53.96	31.38%	—
corpus-10k/ (10,000 stories)	3,956,469	2.9493	19.09	37.43%	−1.039 [−1.078, −1.001]
root, 1,000 stories (max-accuracy)	3,956,469	3.2725	26.38	33.50%	−0.716 [−0.746, −0.686]
root, 1,000 stories (min-CE)	3,956,469	3.2802	26.58	33.37%	−0.708 [−0.738, −0.679]

How much of this is the low rank, and how much is the data?

We tested that directly, by running the same 2×2 at a matched 16,800-update budget:

audit CE	1,000 stories	10,000 stories
full readout (51.3M)	4.9125	3.2619
rank 64 (3.96M)	3.2802	2.9493
readout gap	1.6323	0.3126

80.8% of the low-rank advantage was a small-data artifact. Ten times the text helps the full readout 5.0× more than the low-rank one — the signature of a head that was memorising. A residual 0.313 nats survives, but neither 10k arm converged (1.83 passes), so that remainder may be a convergence-rate difference rather than a quality one. The full-readout control is published at fly-wordbrain-fullreadout-10k.

Read this carefully. The comparison against the released reference is not controlled: its training stories, trainer and tokenizer-fitting population are unknown, and ours were trained on our own 1,000 TinyStories under a reconstructed recipe. The controlled result is internal — holding everything else fixed and changing only the readout parameterization improved held-out CE by 1.63 nats while removing 12.97× the parameters. That points at readout overfitting, not at anatomy — and a randomised-graph control has since confirmed it. Retraining this recipe on the same connectome with every edge rewired at random, holding in-degree, out-degree and every synaptic weight fixed, costs 0.0100 nats [+0.0026, +0.0175]: the fly's actual wiring is worth about 1% of the 1.039 nats this model gains over the reference. A zero-edge control, which would bound what the recurrence contributes, has not been run.

It is also not fluent. It rambles about Lily and Tom, loses story premises, and sometimes reproduces training phrases verbatim. Low cross-entropy on TinyStories is not writing.

Files

Path	What
corpus-10k/min-ce.safetensors	Recommended. 10,000 stories, update 16,800
corpus-10k/max-accuracy.safetensors	Same run, accuracy selector
corpus-10k/web/ + manifest.json	Browser arrays and receipts for the 10k weights
max-accuracy.safetensors	1,000 stories, update 15,552
min-ce.safetensors	1,000 stories, update 16,600
web/*.f32 + web/manifest.json	Raw float32 arrays the browser demo streams, with a golden trace
manifest.json	Architecture, selector receipts, frozen-graph digests, held-out scores
experiment-config.json	The committed experiment specification

The 10k weights are budget-capped, not trained to convergence: 16,800 updates is 1.83 passes over their corpus. The trainer records them as debug_stopped with debug: true, which flags the cap rather than a failure.

Two selectors are kept as separate artifacts and never merged into one "best" checkpoint. They were selected on different criteria and they are different weights.

The connectome is not duplicated here — it lives once in fly-connectome-49k (the central-brain subset this model runs on; the complete 166,700-neuron graph is fly-connectome-malecns-166k). manifest.json records the frozen-buffer SHA-256 digests that identify exactly which graph these weights were trained on; the dataset manifest carries the same digests.

Run it

import json, numpy as np, torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from scipy.sparse import csr_matrix

weights = snapshot_download('fernandofernandes/fly-wordbrain-rank64')
graph = snapshot_download('fernandofernandes/fly-connectome-49k', repo_type='dataset')

p = load_file(f'{weights}/max-accuracy.safetensors')
offsets = np.fromfile(f'{graph}/graph/edges_offsets.i32', dtype=np.int32)
source = np.fromfile(f'{graph}/graph/edges_source.u16', dtype=np.uint16).astype(np.int32)
value = np.fromfile(f'{graph}/graph/edges_weight.f32', dtype=np.float32)
W = csr_matrix((value, source, offsets), shape=(49393, 49393))
in_index = np.fromfile(f'{graph}/interface/in_index.i32', dtype=np.int32)

The full step function is ~30 lines; see scripts/export_web_model.py, whose NumPy implementation is checked against PyTorch token-for-token, and space/src/engine.js for the JavaScript version that reproduces the same golden trace exactly.

Training

From scratch, seed 42, AdamW, batch 8, TBPTT 32, gradient clipping 1.0, on an Apple M3 Max through custom Metal sparse kernels contributed upstream to ConnecTorch. Stopped at 14 completed epochs / 16,800 updates under a plateau rule — the planned 44-epoch schedule was not completed, so these are accepted early stops rather than finished runs.

TinyStories (CDLA-Sharing-1.0) is not redistributed; the repository carries story identifiers and hashes so the exact splits can be rebuilt from the upstream dataset.

Licence

Weights CC BY 4.0, since they derive from the MaleCNS v1.0 connectome — FlyEM / HHMI Janelia Research Campus, University of Cambridge, MRC Laboratory of Molecular Biology, and Google Research. Architecture and tokenizer after ngxson/fly-llm-hf (CC BY 4.0). Our code is MIT.

A smaller sibling at 2,343,125 parameters: fly-wordbrain-rank32.