nanofly-decoder-en / README.md
igorktech's picture
nanoFLY decoder: MaleCNS v1.0 connectome as a frozen recurrent layer
23cc7a6
|
Raw History Blame Contribute Delete
6.51 kB
---
license: cc-by-4.0
language:
- en
library_name: transformers
pipeline_tag: text-generation
datasets:
- roneneldan/TinyStories
tags:
- connectome
- reservoir-computing
- echo-state-network
- fruit-fly
- drosophila
- malecns
- tinystories
- custom_code
---
# nanofly-decoder-en
![MaleCNS soma cloud coloured by this model's neuron states](connectome.png)
<sub>43,993 neurons at their measured MaleCNS v1.0 coordinates, coloured by this checkpoint's state at one tick while writing *"…there was a girl named Lily."* Orange excited, blue inhibited, grey at rest. Frontal view; the optic lobes flank the central brain.</sub>
A language model whose recurrent layer is the measured wiring of a fruit fly. The connectome is a frozen [echo state network](https://en.wikipedia.org/wiki/Echo_state_network) reservoir β€” **no synapse is trained**. Only the input projection, per-neuron gain/bias/leak, one global scale and the readout learn.
- **Code:** [github.com/igorktech/nanoFLY](https://github.com/igorktech/nanoFLY)
- **Russian sibling:** [`igorktech/nanofly-decoder-ru`](https://huggingface.co/igorktech/nanofly-decoder-ru)
## Usage
```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
repo = "igorktech/nanofly-decoder-en"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True).eval()
model = model.to("cuda" if torch.cuda.is_available() else "cpu")
ids = tok("Once upon a time, there was a", return_tensors="pt").input_ids
ids = torch.cat([torch.tensor([[model.config.bos_token_id]]), ids], dim=1).to(model.device)
out = model.generate(ids, max_new_tokens=80, do_sample=True, top_k=50, temperature=0.7)
print(tok.decode(out[0], skip_special_tokens=True))
```
- Prepend BOS: every training example started with it.
- Greedy: `do_sample=False`, and drop `top_k` or `transformers` warns.
- Beam search and assisted generation are unsupported (stateful model). Greedy, sampling, top-k, top-p work.
- ~960 forward passes/s on an RTX 5080, ~10/s on a laptop CPU.
## Architecture
| | |
|---|---|
| connectome | MaleCNS v1.0 central brain β€” `cb_sensory`, `cb_intrinsic`, `visual_projection`, `descending_neuron`, `ascending_neuron` |
| neurons / edges | 49,393 / 9,055,280 signed (623,728 dropped: modulatory or unknown transmitter) |
| edge weight | sign of the presynaptic transmitter Γ— synapse count, rows normalised to unit absolute weight. ACh +1; GABA, Glu, His βˆ’1; others 0 |
| token input | 11,434 sensory-facing neurons, 8-slot delay line (slot *j* gets token *tβˆ’j*). No attention, no positional encoding |
| held out | the 2,635 ORNs stay out of the token input, so the encoder-decoder variant can start from these weights |
| dynamics | `x ← (1βˆ’a)Β·x + aΒ·tanh(ρ·gΒ·(Wx) + u + b)`, 2 ticks per token; `a` learned per neuron (init 0.5), `ρ` learned global (init 1.0 β†’ **4.26**) |
| readout | all 49,393 states β†’ `Linear(49393β†’256)` β†’ `LayerNorm` β†’ `Linear(256β†’1024)` |
| trainable | **16.25M** β€” readout 12.90M, input projection 2.93M, embedding 0.26M, per-neuron scalars 0.15M |
## Training
| | |
|---|---|
| data | TinyStories, 98,024 stories / 1,976 held out, 28,417,062 tokens, byte-level BPE vocab 1,024 |
| objective | next-token cross entropy, truncated BPTT over 32-token windows, state carried across windows |
| optimiser | AdamW β€” body 2e-3 (no decay), readout 5e-4 (decay 0.01), warmup 200 then cosine to 10%, clip 1.0 |
| schedule | 3 epochs, 22,115 updates, batch 128 |
| hardware | 1 Γ— RTX 5080, 1.74 h at ~13,700 tok/s |
## Evaluation
Held-out loss. Identical data, recipe, seed and budget; only the recurrent layer differs:
| recurrent layer | val loss | ppl | trainable |
|---|---|---|---|
| degree-matched shuffled wiring, frozen | 1.979 | 7.2 | 16.25M |
| **this model** β€” real connectome, frozen | **1.933** | **6.9** | 16.25M |
| real connectome, synapse strengths trained | 1.913 | 6.8 | 25.30M |
The control keeps every neuron's in- and out-degree, the transmitter signs and Dale's law, and randomises only which neuron pairs with which. Two results, each holding at all ten intermediate evaluations:
- **The fly's specific wiring is worth 0.046 nats** over a random graph with its degrees β€” and that margin grows through training (0.031 β†’ 0.046).
- **Making all 9,055,280 synapse strengths trainable is worth 0.020 nats** on top β€” and that margin shrinks through training (0.032 β†’ 0.020), at 4Γ— the compute. It buys faster convergence more than a higher ceiling. Trained strengths stay close to the measured ones (Pearson r = 0.94).
One seed per condition, so treat the margins as indicative rather than significant. Validation was still improving at the end in all three runs; none is converged.
Greedy samples, prompt in bold:
> **Once upon a time, there was a** little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, red ball that was very pretty. She wanted to play with it, but it was too high for her to reach.
> **Tom and his dog** were best friends. They liked to play in the park. One day, they saw a big dog with a big bag.
## Limitations
- 16M trainable parameters over 85M tokens of children's stories. It writes about Lily and Tom and nothing else.
- 8-token delay line plus a short leaky recurrent memory. Loss stops improving with context by roughly position 32; names and objects drift within a paragraph.
- A `tanh` rate neuron is not a spiking model: no spikes, no synaptic delays, no neuromodulation β€” modulatory edges are removed outright.
- Central brain only; the optic lobes and ventral nerve cord of the 166,700-neuron CNS are absent.
- Synapse count is a proxy for strength, and rows are normalised. Neither is measured physiology β€” see the third row of the evaluation table for what happens when the strengths are fit to the task instead.
## Credits
- **Connectome:** MaleCNS v1.0 β€” FlyEM / HHMI Janelia, University of Cambridge, MRC LMB, Google Research. CC BY 4.0. The published buffers derive from that release; keep the attribution when redistributing.
- **Transmitter signs:** Shiu et al., *Nature* 2024.
- **Connectome as reservoir:** Costi, Hadjiivanov, Dold, Hale, Izzo, 2025.
- **Prior art:** [`ngxson/fly-llm-hf`](https://huggingface.co/ngxson/fly-llm-hf), whose graph subset this reproduces.
- **Data:** Eldan and Li, *TinyStories*, 2023.
Weights CC BY 4.0, matching the connectome. Modeling code Apache-2.0.