Whole fly nervous system: 160,514 neurons, 10.4M trainable synapses, best val 1.5957

full-cns/README.md

@@ -0,0 +1,230 @@
+---
+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.
+
+**Trained.** Condition `real`, seed 1, step 15375, validation loss 1.5957 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.
+
+| | |
+|---|---|
+| Source | MaleCNS v1.0 flat connectome (`gs://flyem-male-cns/v1.0/connectome-data/flat-connectome/`) |
+| Candidate pool | central brain: `superclass` starting with `cb_` (164,587 neurons) |
+| Minimum synapses per connection | 3 (engineering choice, not a biological claim) |
+| Extraction | largest SCC → largest directed (k,k)-core with ≥ target nodes (k = 0) → trim by weighted degree |
+| Neurons used | 160,514 |
+| Directed connections used | 10,402,842 |
+| Synaptic contacts represented | 102,352,176 |
+| Largest SCC fraction | 1.0 |
+| Reciprocal pairs | 1,006,326 |
+| Input / output neurons | top 2048 by out-degree / top 4096 by in-degree |
+| Input→output shortest path (median / p90 / max hops) | 1.0 / 1.0 / 1.0 |
+| Graph hash | `d53f26f9bf6f7b8c9944e28fdd4721d2c0a237e7534f9160977c70e6c81ac6c2` |
+
+`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, 10402842) | int32 | 83.22 MB |
+| `graph.synapse_count` | (10402842,) | int32 | 41.61 MB |
+| `graph.node_id` | (160514,) | int64 | 1.28 MB |
+| `graph.input_nodes` | (2048,) | int64 | 0.02 MB |
+| `graph.output_nodes` | (4096,) | int64 | 0.03 MB |
+| `recurrent.edge_values` | (10402842,) | bfloat16 | 20.81 MB |
+| `recurrent.bias` | (160514,) | bfloat16 | 0.32 MB |
+| `recurrent.raw_leak` | (160514,) | bfloat16 | 0.32 MB |
+| `embed.weight` | (65, 32) | bfloat16 | 0.00 MB |
+| `input_proj.weight` | (2048, 32) | bfloat16 | 0.13 MB |
+| `input_proj.bias` | (2048,) | bfloat16 | 0.00 MB |
+| `lm_head.weight` | (65, 4096) | bfloat16 | 0.53 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 → 2048 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))
+4096 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="full-cns", 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)  # the 5k control
+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="full-cns", 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")
+```
+
+## Result: the whole nervous system
+
+This checkpoint is the **entire traced fly nervous system**: 160,514 neurons and 10,402,842 trainable synaptic
+connections, covering the optic lobes, central brain and ventral nerve cord. Trained for 16,700 steps at global
+batch 192 across six GPUs (205M characters, about 200 epochs of Tiny Shakespeare).
+
+Best validation loss **1.5957** nats/char, against a bigram floor of 2.482 on the same split.
+
+**Scale did not improve accuracy.** The 5,000-neuron central-brain subgraph reaches 1.578 with one eighteenth of
+the neurons. At a 64-character context on a 1.1M-character corpus the binding constraint is the context window
+and the data, not model capacity, and the larger graph overfits sooner (train loss 1.333 against val 1.596).
+Text quality nonetheless reads better by eye, with correctly spelled rare character names and more grammatical
+lines.
+
+This is **one seed with one control**, not a five-seed claim. The wiring comparison that the pre-registered rule
+in the project's spec governs was run at 5,000 neurons and is reported in this repository's main model card.
+
+## 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).

full-cns/config.json

@@ -0,0 +1,42 @@
+{
+  "model_type": "flygpt",
+  "architectures": [
+    "FlyGPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_flygpt.FlyGPTConfig",
+    "AutoModelForCausalLM": "modeling_flygpt.FlyGPTForCausalLM"
+  },
+  "vocab_size": 65,
+  "num_neurons": 160514,
+  "num_edges": 10402842,
+  "embedding_dim": 32,
+  "num_input_neurons": 2048,
+  "num_output_neurons": 4096,
+  "microsteps": 2,
+  "activation": "tanh",
+  "learned_leak": true,
+  "leak_init": 0.5,
+  "degree_normalization": true,
+  "init_scale": 0.65,
+  "dtype": "bfloat16",
+  "training_steps": 16700,
+  "graph": {
+    "graph_name": "cns_full",
+    "condition": "real",
+    "control_seed": 1,
+    "source": "malecns-v1.0",
+    "region_filter": "whole_cns",
+    "min_synapses": 3,
+    "hash": "d53f26f9bf6f7b8c9944e28fdd4721d2c0a237e7534f9160977c70e6c81ac6c2"
+  },
+  "training_state": {
+    "status": "trained",
+    "condition": "real",
+    "seed": 1,
+    "step": 15375,
+    "val_loss": 1.5957297205924987,
+    "checkpoint": "checkpoints/flygpt-v0-fullcns/cns_full/real_seed1.pt",
+    "project": "flygpt-v0-fullcns"
+  }
+}

full-cns/configuration_flygpt.py

@@ -0,0 +1,40 @@
+"""FlyGPT configuration for Hugging Face `transformers` (loaded with trust_remote_code=True)."""
+from transformers import PretrainedConfig
+
+
+class FlyGPTConfig(PretrainedConfig):
+    model_type = "flygpt"
+
+    def __init__(
+        self,
+        vocab_size: int = 65,
+        num_neurons: int = 0,
+        num_edges: int = 0,
+        embedding_dim: int = 32,
+        num_input_neurons: int = 256,
+        num_output_neurons: int = 512,
+        microsteps: int = 2,
+        activation: str = "tanh",
+        learned_leak: bool = True,
+        leak_init: float = 0.5,
+        degree_normalization: bool = True,
+        init_scale: float = 0.1,
+        graph: dict | None = None,
+        training_state: dict | None = None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.num_neurons = num_neurons
+        self.num_edges = num_edges
+        self.embedding_dim = embedding_dim
+        self.num_input_neurons = num_input_neurons
+        self.num_output_neurons = num_output_neurons
+        self.microsteps = microsteps
+        self.activation = activation
+        self.learned_leak = learned_leak
+        self.leak_init = leak_init
+        self.degree_normalization = degree_normalization
+        self.init_scale = init_scale
+        self.graph = graph or {}                    # source, hash, region filter, stats, provenance
+        self.training_state = training_state or {}  # "init" or the run that produced the values
+        super().__init__(**kwargs)

full-cns/generation_config.json

@@ -0,0 +1,5 @@
+{
+  "do_sample": true,
+  "temperature": 0.8,
+  "max_new_tokens": 300
+}

full-cns/graph_metadata.json

@@ -0,0 +1,78 @@
+{
+  "graph_name": "cns_full",
+  "condition": "real",
+  "control_seed": 1,
+  "source": "malecns-v1.0",
+  "region_filter": "whole_cns",
+  "min_synapses": 3,
+  "selector": "directed_core",
+  "target_neurons": 200000,
+  "hash": "d53f26f9bf6f7b8c9944e28fdd4721d2c0a237e7534f9160977c70e6c81ac6c2",
+  "stats": {
+    "candidate_pool": 164587,
+    "candidate_edges": 25563096,
+    "min_synapses": 3,
+    "target_neurons": 200000,
+    "edges_after_threshold": 10510984,
+    "scc_after_threshold": 160514,
+    "note": "SCC no larger than target; no pruning or trimming applied",
+    "core_k": 0,
+    "final_neurons": 160514,
+    "final_edges": 10402842,
+    "source": "malecns-v1.0",
+    "n_neurons": 160514,
+    "n_edges": 10402842,
+    "synaptic_contacts": 102352176.0,
+    "largest_scc_fraction": 1.0,
+    "reciprocal_pairs": 1006326,
+    "in_degree": {
+      "mean": 64.80956178277285,
+      "min": 1,
+      "max": 9193
+    },
+    "out_degree": {
+      "mean": 64.80956178277285,
+      "min": 1,
+      "max": 7749
+    },
+    "density": 0.00040376518900508274,
+    "hash": "d53f26f9bf6f7b8c9944e28fdd4721d2c0a237e7534f9160977c70e6c81ac6c2",
+    "region_breakdown": {
+      "optic_lobe": 91082,
+      "central_brain": 36625,
+      "vnc": 19300,
+      "other": 13454,
+      "unannotated": 53
+    }
+  },
+  "diagnostics": {
+    "provenance": {
+      "control": "real"
+    },
+    "io": {
+      "n_input": 2048,
+      "n_output": 4096,
+      "input_rule": "top_out_degree",
+      "output_rule": "top_in_degree",
+      "overlap": 1381
+    },
+    "diagnostics": {
+      "n_neurons": 160514,
+      "n_edges": 10402842,
+      "largest_scc_fraction": 1.0,
+      "reciprocal_pairs": 1006326,
+      "io_reachable_fraction": 1.0,
+      "io_path_median": 1.0,
+      "io_path_p90": 1.0,
+      "io_path_max": 1.0
+    },
+    "gate": {
+      "microsteps": 2,
+      "max_p90_hops": 6,
+      "min_reachable_fraction": 0.95,
+      "reachable_ok": true,
+      "p90_ok": true,
+      "passed": true
+    }
+  }
+}

full-cns/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63dc73657502dad71a475dc319b0acf38a6d6b89da57954141f08fd728f7441a
+size 148288294

full-cns/modeling_flygpt.py

@@ -0,0 +1,199 @@
+"""FlyGPT: a character-level language model whose recurrent core is a real subgraph of the
+fruit-fly connectome (MaleCNS v1.0). Hugging Face `transformers` implementation; self-contained.
+
+Dynamics (one scalar state per neuron, plan.md §7 of the FlyGPT spec):
+
+    proposal_i = tanh( sum_j W_ij h_j / sqrt(in_degree_i) + external_input_i + bias_i )
+    h_i_new    = (1 - leak_i) * h_i + leak_i * proposal_i
+
+The connectome is stored in `model.safetensors` as integer tensors (`graph.*`); only the learned
+per-edge values and the adapters are floating point (bf16 on disk). The sparse recurrent matmul runs
+in fp32: through the fused kernels of the `connectome-kernels` package when it is installed and a
+CUDA device is used (training speed), else through torch.sparse COO (rows = destination). Both give
+the same logits and gradients.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.utils import ModelOutput
+
+from .configuration_flygpt import FlyGPTConfig
+
+
+@dataclass
+class FlyGPTOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    state: Optional[torch.FloatTensor] = None   # [B, N] neuron states after the last character
+
+
+class FlyGraph(nn.Module):
+    """The anatomy. Integer buffers only; never trained."""
+
+    def __init__(self, num_neurons: int, num_edges: int, num_input: int, num_output: int):
+        super().__init__()
+        self.register_buffer("edge_index", torch.zeros(2, num_edges, dtype=torch.int32))   # [source, destination]
+        self.register_buffer("synapse_count", torch.zeros(num_edges, dtype=torch.int32))   # MaleCNS synaptic contacts
+        self.register_buffer("node_id", torch.zeros(num_neurons, dtype=torch.int64))       # MaleCNS body ids
+        self.register_buffer("input_nodes", torch.zeros(num_input, dtype=torch.int64))
+        self.register_buffer("output_nodes", torch.zeros(num_output, dtype=torch.int64))
+
+
+class FlyRecurrentCore(nn.Module):
+    """The learned state: one value per real edge, plus per-neuron bias and leak."""
+
+    def __init__(self, num_neurons: int, num_edges: int, leak_init: float, learned_leak: bool):
+        super().__init__()
+        self.edge_values = nn.Parameter(torch.zeros(num_edges))
+        self.bias = nn.Parameter(torch.zeros(num_neurons))
+        self.raw_leak = nn.Parameter(torch.full((num_neurons,), math.log(leak_init / (1 - leak_init))),
+                                     requires_grad=learned_leak)
+
+
+class FlyGPTPreTrainedModel(PreTrainedModel):
+    config_class = FlyGPTConfig
+    base_model_prefix = "flygpt"
+    _is_stateful = True
+    _supports_cache_class = False
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        if isinstance(module, FlyRecurrentCore):
+            nn.init.normal_(module.edge_values, std=self.config.init_scale)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=1.0)
+
+
+class FlyGPTForCausalLM(FlyGPTPreTrainedModel, GenerationMixin):
+    def __init__(self, config: FlyGPTConfig):
+        super().__init__(config)
+        c = config
+        self.graph = FlyGraph(c.num_neurons, c.num_edges, c.num_input_neurons, c.num_output_neurons)
+        self.recurrent = FlyRecurrentCore(c.num_neurons, c.num_edges, c.leak_init, c.learned_leak)
+        self.embed = nn.Embedding(c.vocab_size, c.embedding_dim)
+        self.input_proj = nn.Linear(c.embedding_dim, c.num_input_neurons)
+        self.lm_head = nn.Linear(c.num_output_neurons, c.vocab_size)
+        self.post_init()
+
+    # ---- sparse recurrent matrix -------------------------------------------------------------
+    @property
+    def num_neurons(self) -> int:
+        return self.config.num_neurons
+
+    def edge_scale(self) -> torch.Tensor:
+        """1/sqrt(in_degree) per edge (degree normalization), or ones."""
+        dst = self.graph.edge_index[1].long()
+        if not self.config.degree_normalization:
+            return torch.ones_like(dst, dtype=torch.float32)
+        in_deg = torch.bincount(dst, minlength=self.num_neurons).clamp(min=1).float()
+        return 1.0 / in_deg[dst].sqrt()
+
+    def sparse_weight(self) -> torch.Tensor:
+        src, dst = self.graph.edge_index[0].long(), self.graph.edge_index[1].long()
+        values = self.recurrent.edge_values.float() * self.edge_scale()
+        return torch.sparse_coo_tensor(torch.stack([dst, src]), values, (self.num_neurons, self.num_neurons))
+
+    def dense_weight(self) -> torch.Tensor:
+        """Convenience for analysis; [N, N] with rows = destination. Never used in the forward pass."""
+        return self.sparse_weight().to_dense()
+
+    @property
+    def leak(self) -> torch.Tensor:
+        return torch.sigmoid(self.recurrent.raw_leak.float())
+
+    # ---- dynamics ------------------------------------------------------------------------------
+    def init_state(self, batch: int, device=None) -> torch.Tensor:
+        return torch.zeros(batch, self.num_neurons, device=device or self.recurrent.edge_values.device)
+
+    def drive(self, x: torch.Tensor) -> torch.Tensor:
+        d = torch.zeros(x.shape[0], self.num_neurons, device=x.device, dtype=torch.float32)
+        d[:, self.graph.input_nodes] = self.input_proj(self.embed(x)).float()
+        return d
+
+    def step(self, state: torch.Tensor, x: torch.Tensor, W: torch.Tensor | None = None) -> torch.Tensor:
+        W = self.sparse_weight() if W is None else W
+        drive, leak, bias = self.drive(x), self.leak, self.recurrent.bias.float()
+        for _ in range(self.config.microsteps):
+            incoming = torch.sparse.mm(W, state.float().T).T
+            proposal = torch.tanh(incoming + drive + bias)
+            state = (1 - leak) * state + leak * proposal
+        return state
+
+    def logits_from_state(self, state: torch.Tensor) -> torch.Tensor:
+        return self.lm_head(state[:, self.graph.output_nodes].to(self.lm_head.weight.dtype)).float()
+
+    # ---- fused CUDA path via the connectome-kernels package (optional, used for training) ------------
+    def _fused_graph(self):
+        from connectome_kernels import SparseGraph
+        dev = self.graph.edge_index.device
+        if getattr(self, "_fg", None) is None or self._fg_device != dev:
+            self._fg = SparseGraph(self.graph.edge_index[0].long(), self.graph.edge_index[1].long(), self.num_neurons,
+                                   self.graph.input_nodes)
+            self._fg_device = dev
+        return self._fg
+
+    def _fused_available(self, device) -> bool:
+        if device.type != "cuda":
+            return False
+        if not hasattr(self, "_fused_ok"):
+            try:
+                from connectome_kernels import available
+                self._fused_ok = available()
+            except Exception:
+                self._fused_ok = False
+        return self._fused_ok
+
+    def _forward_fused(self, input_ids, state):
+        from connectome_kernels import sparse_recurrence
+        drives = self.input_proj(self.embed(input_ids)).float().permute(1, 2, 0).contiguous()       # [T, n_in, B]
+        vals = self.recurrent.edge_values.float() * self.edge_scale()
+        out = sparse_recurrence(vals, self.leak, self.recurrent.bias.float(), drives, state,
+                                self._fused_graph(), self.config.microsteps)                         # [T, B, N]
+        logits = self.lm_head(out[:, :, self.graph.output_nodes].to(self.lm_head.weight.dtype)).float().permute(1, 0, 2)
+        return logits, out[-1]
+
+    def forward(self, input_ids: torch.LongTensor, state: Optional[torch.Tensor] = None,
+                labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None,
+                return_dict: Optional[bool] = None, **kwargs) -> FlyGPTOutput:
+        B, T = input_ids.shape
+        state = self.init_state(B, input_ids.device) if state is None else state
+        if self._fused_available(input_ids.device):
+            logits, state = self._forward_fused(input_ids, state)
+        else:
+            W = self.sparse_weight()
+            outs = []
+            for t in range(T):
+                state = self.step(state, input_ids[:, t], W)
+                outs.append(self.logits_from_state(state))
+            logits = torch.stack(outs, 1)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.shape[-1]), labels[:, 1:].reshape(-1))
+        return FlyGPTOutput(loss=loss, logits=logits, state=state)
+
+    # ---- generation: carry the neuron state instead of a KV cache ------------------------------
+    @classmethod
+    def _supports_default_dynamic_cache(cls) -> bool:
+        return False  # stateful recurrent model: no KV cache, the neuron state is carried in `state`
+
+    def prepare_inputs_for_generation(self, input_ids, state=None, **kwargs):
+        if state is not None:
+            input_ids = input_ids[:, -1:]
+        return {"input_ids": input_ids, "state": state}
+
+    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, **kwargs):
+        model_kwargs["state"] = outputs.state
+        return model_kwargs

full-cns/tokenizer.json

@@ -0,0 +1,90 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": null,
+  "added_tokens": [],
+  "normalizer": null,
+  "pre_tokenizer": {
+    "type": "Split",
+    "pattern": {
+      "Regex": "[\\s\\S]"
+    },
+    "behavior": "Isolated",
+    "invert": false
+  },
+  "post_processor": null,
+  "decoder": {
+    "type": "Fuse"
+  },
+  "model": {
+    "type": "WordLevel",
+    "vocab": {
+      "\n": 0,
+      " ": 1,
+      "!": 2,
+      "$": 3,
+      "&": 4,
+      "'": 5,
+      ",": 6,
+      "-": 7,
+      ".": 8,
+      "3": 9,
+      ":": 10,
+      ";": 11,
+      "?": 12,
+      "A": 13,
+      "B": 14,
+      "C": 15,
+      "D": 16,
+      "E": 17,
+      "F": 18,
+      "G": 19,
+      "H": 20,
+      "I": 21,
+      "J": 22,
+      "K": 23,
+      "L": 24,
+      "M": 25,
+      "N": 26,
+      "O": 27,
+      "P": 28,
+      "Q": 29,
+      "R": 30,
+      "S": 31,
+      "T": 32,
+      "U": 33,
+      "V": 34,
+      "W": 35,
+      "X": 36,
+      "Y": 37,
+      "Z": 38,
+      "a": 39,
+      "b": 40,
+      "c": 41,
+      "d": 42,
+      "e": 43,
+      "f": 44,
+      "g": 45,
+      "h": 46,
+      "i": 47,
+      "j": 48,
+      "k": 49,
+      "l": 50,
+      "m": 51,
+      "n": 52,
+      "o": 53,
+      "p": 54,
+      "q": 55,
+      "r": 56,
+      "s": 57,
+      "t": 58,
+      "u": 59,
+      "v": 60,
+      "w": 61,
+      "x": 62,
+      "y": 63,
+      "z": 64
+    },
+    "unk_token": "<unk>"
+  }
+}

full-cns/tokenizer_config.json

@@ -0,0 +1,6 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 10000000,
+  "clean_up_tokenization_spaces": false,
+  "add_prefix_space": false
+}