FlyGPT

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