FlyGPT

A character-level language model whose recurrent architecture is a real subgraph of the fruit-fly brain connectome (MaleCNS v1.0). Unlike the earlier frozen-reservoir approach in 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, 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.

Trained. Condition real, seed 1, step 94000, validation loss 1.5778 nats/char on the fixed Tiny Shakespeare split.

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.

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

What is in model.safetensors

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

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

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("QuixiAI/FlyGPT")
model = AutoModelForCausalLM.from_pretrained("QuixiAI/FlyGPT", 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 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.

# 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", trust_remote_code=True, dtype=torch.float32).cuda()
# start from the untrained initialization instead:  subfolder="init"

text = requests.ge