init/README.md

---
license: cc-by-4.0
language: [en]
library_name: transformers
pipeline_tag: text-generation
base_model: QuixiAI/MaleCNS
base_model_relation: finetune
datasets: [karpathy/tiny_shakespeare]
tags: [connectome, fruit-fly, drosophila, malecns, recurrent, sparse, tiny-shakespeare, custom_code]
---

# QuixiAI/FlyGPT

A character-level language model whose recurrent architecture is **a real subgraph of the fruit-fly brain
connectome** ([MaleCNS v1.0](https://male-cns.janelia.org/)). Unlike the earlier frozen-reservoir approach in
[ngxson/fly-llm-hf](https://huggingface.co/ngxson/fly-llm-hf), which keeps the connectome's synaptic weights fixed and
trains only the projections and readout, FlyGPT **trains one value per real synaptic connection with gradient
descent** while keeping the fly's edge topology fixed, and compares the result against the same neurons with
degree-preserving scrambled connections across paired seeds.

**Base model: [QuixiAI/MaleCNS](https://huggingface.co/QuixiAI/MaleCNS)**, the lossless packaging of the MaleCNS v1.0
connectivity tables. FlyGPT's graph is extracted from it deterministically (`build_graph.py`, revision pinned in
`data/fly/build_edges.py`); `graph.node_id` and `graph.synapse_count` map every edge back to that repository.

This checkpoint's wiring is the original MaleCNS wiring.

**Untrained initialization.** The recurrent edge values and adapters are the exact random initialization that FlyGPT's `train.py` starts from for seed 1. Text generated from this checkpoint is noise by design.

This is not a biological simulation of a living fly. The "weights" in the MaleCNS release are anatomical synapse
counts; they are stored here as `graph.synapse_count` and are **not** the model's parameters.

## The graph

Every number below is produced by FlyGPT's extraction script (`build_graph.py`), not typed by hand.

| | |
|---|---|
| Source | MaleCNS v1.0 flat connectome (`gs://flyem-male-cns/v1.0/connectome-data/flat-connectome/`) |
| Candidate pool | central brain: `superclass` starting with `cb_` (37,108 neurons) |
| Minimum synapses per connection | 3 (engineering choice, not a biological claim) |
| Extraction | largest SCC → largest directed (k,k)-core with ≥ target nodes (k = 40) → trim by weighted degree |
| Neurons used | 5,000 |
| Directed connections used | 524,324 |
| Synaptic contacts represented | 8,300,915 |
| Largest SCC fraction | 1.0 |
| Reciprocal pairs | 93,055 |
| Input / output neurons | top 256 by out-degree / top 512 by in-degree |
| Input→output shortest path (median / p90 / max hops) | 1.0 / 1.0 / 1.0 |
| Graph hash | `f82b783b7ccb5a354fc4cf3de6de4a98d75029303c55f8faae28ab807828a007` |

`graph.node_id` holds the MaleCNS body ids, so every neuron maps back to the release.

## What is in `model.safetensors`

| tensor | shape | dtype | size |
|---|---|---|---|
| `graph.edge_index` | (2, 524324) | int32 | 4.19 MB |
| `graph.synapse_count` | (524324,) | int32 | 2.10 MB |
| `graph.node_id` | (5000,) | int64 | 0.04 MB |
| `graph.input_nodes` | (256,) | int64 | 0.00 MB |
| `graph.output_nodes` | (512,) | int64 | 0.00 MB |
| `recurrent.edge_values` | (524324,) | bfloat16 | 1.05 MB |
| `recurrent.bias` | (5000,) | bfloat16 | 0.01 MB |
| `recurrent.raw_leak` | (5000,) | bfloat16 | 0.01 MB |
| `embed.weight` | (65, 32) | bfloat16 | 0.00 MB |
| `input_proj.weight` | (256, 32) | bfloat16 | 0.02 MB |
| `input_proj.bias` | (256,) | bfloat16 | 0.00 MB |
| `lm_head.weight` | (65, 512) | bfloat16 | 0.07 MB |
| `lm_head.bias` | (65,) | bfloat16 | 0.00 MB |

`graph.*` is the anatomy (integer, never trained). `recurrent.*`, `embed.*`, `input_proj.*`, `lm_head.*` are the
learned state, stored in bf16. The sparse recurrent matmul is rebuilt in fp32 at runtime (rows = destination,
columns = source), with each incoming edge scaled by `1/sqrt(in_degree)`.

## Dynamics

```text
character → embedding (32) → linear → 256 input neurons
proposal_i = tanh( Σ_j W_ij h_j / sqrt(in_degree_i) + external_input_i + bias_i )
h_i ← (1 − leak_i) h_i + leak_i · proposal_i        (2 microsteps per character, leak_i = sigmoid(raw_leak_i))
512 output neuron states → linear → 65 logits
```

## Inference

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("QuixiAI/FlyGPT")
model = AutoModelForCausalLM.from_pretrained("QuixiAI/FlyGPT", subfolder="init", trust_remote_code=True, dtype=torch.float32)

ids = tok("ROMEO:", return_tensors="pt").input_ids
out = model.generate(ids, max_new_tokens=300, do_sample=True, temperature=0.8)
print(tok.decode(out[0]))

# The degree-preserving scrambled control (same neurons, same degrees, shuffled wiring), for comparison:
scrambled = AutoModelForCausalLM.from_pretrained("QuixiAI/FlyGPT", subfolder="scrambled", trust_remote_code=True, dtype=torch.float32)
print(tok.decode(scrambled.generate(ids, max_new_tokens=300, do_sample=True, temperature=0.8)[0]))

# Neuron activity, for visualization: [1, T, 5000] states after each character, plus MaleCNS body ids
with torch.no_grad():
    states = model(ids).state                                   # [B, N] after the last character
body_ids = model.graph.node_id                                  # index -> MaleCNS body id, for lookup in QuixiAI/MaleCNS
```

The tokenizer is strict: only the 65 characters of Tiny Shakespeare are encodable. `generate()` carries the neuron
state between characters instead of a KV cache.

## Training

The recurrent core has one trainable weight per real synaptic connection. With the
[connectome-kernels](https://github.com/QuixiAI/connectome-kernels) package installed, the model's forward pass
runs on fused CUDA kernels (about 13× faster than `torch.sparse`, identical gradients); without it, it falls back
to `torch.sparse` automatically.

```python
# Fine-tune / continue training FlyGPT on Tiny Shakespeare (character-level).
# pip install transformers safetensors
# pip install --no-build-isolation git+https://github.com/QuixiAI/connectome-kernels   # fused CUDA path, ~13x faster
import requests, torch, torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("QuixiAI/FlyGPT")
model = AutoModelForCausalLM.from_pretrained("QuixiAI/FlyGPT", subfolder="init", trust_remote_code=True, dtype=torch.float32).cuda()
# start from the untrained initialization instead:  subfolder="init"

text = requests.get("https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt").text
data = torch.tensor(tok(text).input_ids)
train, val = data[: int(0.9 * len(data))], data[int(0.9 * len(data)):]   # FlyGPT's fixed 90/10 split

def batch(split, B=32, T=64):
    i = torch.randint(0, len(split) - T - 1, (B,))
    x = torch.stack([split[j : j + T] for j in i]); y = torch.stack([split[j + 1 : j + T + 1] for j in i])
    return x.cuda(), y.cuda()

recurrent = list(model.recurrent.parameters())                         # one weight per real synapse, bias, leak
adapters = [p for n, p in model.named_parameters() if not n.startswith("recurrent.")]
opt = torch.optim.AdamW([{"params": adapters, "lr": 1e-3}, {"params": recurrent, "lr": 3e-4}], weight_decay=0.01)

for step in range(1, 501):
    x, y = batch(train)
    logits = model(x).logits                                           # [B, T, 65]; state resets to zero per window
    loss = F.cross_entropy(logits.reshape(-1, 65), y.reshape(-1))
    opt.zero_grad(set_to_none=True); loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
    if step % 100 == 0:
        with torch.no_grad():
            vx, vy = batch(val); vl = F.cross_entropy(model(vx).logits.reshape(-1, 65), vy.reshape(-1))
        print(f"step {step}  train {loss.item():.3f}  val {vl.item():.3f}")

model.save_pretrained("flygpt-finetuned"); tok.save_pretrained("flygpt-finetuned")
```

## Citation

If you use this model, please cite it, its base model, and the MaleCNS dataset paper.

This model:

```bibtex
@misc{hartford2026flygpt,
  title        = {FlyGPT: a language model whose recurrent architecture is a real subgraph of the fruit-fly connectome},
  author       = {Hartford, Eric},
  year         = {2026},
  publisher    = {Hugging Face},
  howpublished = {\url{https://huggingface.co/QuixiAI/FlyGPT}},
  note         = {Base model: QuixiAI/MaleCNS (MaleCNS v1.0, Berg et al. 2026, CC-BY 4.0). Code: https://github.com/QuixiAI/FlyGPT}
}
```

The base model (lossless connectome packaging):

```bibtex
@misc{hartford2026malecns,
  title        = {QuixiAI/MaleCNS: the MaleCNS v1.0 fruit-fly connectome as lossless Safetensors},
  author       = {Hartford, Eric},
  year         = {2026},
  publisher    = {Hugging Face},
  doi          = {10.57967/hf/10410},
  howpublished = {\url{https://huggingface.co/QuixiAI/MaleCNS}},
  note         = {Repackaging of Berg et al. (2026), CC-BY 4.0}
}
```

The dataset (required by the CC-BY 4.0 license):

```bibtex
@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ł and Marin, Elizabeth C. and Nern, Aljoscha and Preibisch, Stephan and Qiu, Wei and Takemura, Shin-ya and Fragniere, Alexandra M.C. and Champion, Andrew S. and Adjavon, Diane-Yayra and Cook, Michael and Gkantia, Marina and Hayworth, Kenneth J. and Huang, Gary B. and Katz, William T. and Kämpf, Florian and Lu, Zhiyuan and Ordish, Christopher and Paterson, Tyler and Stürner, Tomke and Trautman, Eric T. and Whittle, Catherine R. and Burnett, Laura E. and Hoeller, Judith and Li, Feng and Loesche, Frank and Morris, Billy J. and Pietzsch, Tobias and Pleijzier, Markus W. and Silva, Valeria and Yin, Yijie and Ali, Iris and Badalamente, Griffin and Bates, Alexander Shakeel and Beresford, Rory J. and Bogovic, John and Brooks, Paul and Cachero, Sebastian and Canino, Brandon S. and Chaisrisawatsuk, Bhumpanya and Clements, Jody and Crowe, Arthur and de Haan Vicente, Inês and Dempsey, Georgia and Donà, Erika and Dos Santos, Márcia and Dreher, Marisa and Dunne, Christopher R. and Eichler, Katharina and Finley-May, Samantha and Flynn, Miriam A. and Hameed, Imran and Hopkins, Gary Patrick and Hubbard, Philip M. and Kiassat, Ladann and Kovalyak, Julie and Lauchie, Shirley A. and Leonard, Meghan and Lohff, Alanna and Longden, Kit D. and Maldonado, Charli A. and Moitra, Ilina and Moon, Sung Soo and Mooney, Caroline and Munnelly, Eva J. and Okeoma, Nneoma and Olbris, Donald J. and Pai, Anika and Patel, Birava and Phillips, Emily M. and Plaza, Stephen M. and Richards, Alana and Rivas Salinas, Jennifer and Roberts, Ruairí J.V. and Rogers, Edward M. and Scott, Ashley L. and Scuderi, Louis A. and Seenivasan, Pavithraa and Serratosa Capdevila, Laia and Smith, Claire and Svirskas, Rob and Takemura, Satoko and Tastekin, Ibrahim and Thomson, Alexander and Umayam, Lowell and Walsh, John J. and Whittome, Holly and Xu, C. Shan and Yakal, Emily A. and Yang, Tansy and Zhao, Arthur and George, Reed and Jain, Viren and Jayaraman, Vivek and Korff, Wyatt and Meissner, Geoffrey W. and Romani, Sandro and Funke, Jan and Knecht, Christopher and Saalfeld, Stephan and Scheffer, Louis K. and Waddell, Scott and Card, Gwyneth M. and Ribeiro, Carlos and Reiser, Michael B. and Hess, Harald F. and Rubin, Gerald M. and Jefferis, Gregory S.X.E.},
  journal   = {Cell},
  volume    = {189},
  number    = {18},
  pages     = {5504--5526.e15},
  year      = {2026},
  month     = sep,
  publisher = {Elsevier},
  doi       = {10.1016/j.cell.2026.08.015},
  url       = {https://doi.org/10.1016/j.cell.2026.08.015},
  note      = {Preprint: bioRxiv 10.1101/2025.10.09.680999. Data: MaleCNS v1.0, CC-BY 4.0, https://male-cns.janelia.org}
}
```

## License

The connectome is released under CC-BY 4.0 by the FlyEM Project Team (HHMI Janelia), the University of Cambridge,
the MRC Laboratory of Molecular Biology, and Google Research. This checkpoint is a derivative and carries the same
license.

Prior art: [ngxson/fly-llm-hf](https://huggingface.co/ngxson/fly-llm-hf) (frozen MaleCNS reservoir LM) and
[eob/gpt-fly](https://huggingface.co/eob/gpt-fly) (FlyWire-masked GPT-2). Code and experiment:
[github.com/QuixiAI/FlyGPT](https://github.com/QuixiAI/FlyGPT).